Support for classes that extend DIContainer with new option

README.md

@@ -272,11 +272,12 @@ be extremely fast.
 
 You can pass in a few options to DI-Compiler via command line options:
 
-| Environment Variable        | Description                                                                                                                 |
-| --------------------------- | --------------------------------------------------------------------------------------------------------------------------- |
-| `DI_COMPILER_TSCONFIG_PATH` | The path to the `tsconfig.json` file to use                                                                                 |
-| `DI_COMPILER_IDENTIFIER`    | A comma-separated list of identifiers that should be considered instances of DIContainer when transforming the source files |
-| `DI_COMPILER_DISABLE_CACHE` | If set, no disk caching will be used.                                                                                       |
+| Environment Variable        | Description                                                                                                                   |
+| --------------------------- | ----------------------------------------------------------------------------------------------------------------------------- |
+| `DI_COMPILER_TSCONFIG_PATH` | The path to the `tsconfig.json` file to use                                                                                   |
+| `DI_COMPILER_IDENTIFIER`    | A comma-separated list of identifiers that should be considered instances of `DIContainer` when transforming the source files |
+| `DI_COMPILER_CLASS_NAME`    | A comma-separated list of additional class names for `DIContainer`, for custom extensions that inherit the class.             |
+| `DI_COMPILER_DISABLE_CACHE` | If set, no disk caching will be used.                                                                                         |
 
 Alternatively, you can add a `di` property to your `tsconfig` where you can customize its behavior without setting environment variables:
 
@@ -285,6 +286,7 @@ Alternatively, you can add a `di` property to your `tsconfig` where you can cust
 {
 	"di": {
 		"identifier": "container",
+		"diClassName": ["CustomDIContainer", "AnotherDIContainer"],
 		"disableCache": false
 	},
 	"compilerOptions": {

src/loader/shared.ts

@@ -4,10 +4,13 @@ import {FileCache} from "../transformer/cache.js";
 import type {TransformOptions, TransformResult} from "../transformer/transform-options.js";
 import type {TS} from "../type/type.js";
 import {booleanize} from "../util/util.js";
+import { DI_CONTAINER_NAME } from '../transformer/constant.js';
 
 export const ENV_VARIABLE_TSCONFIG_PATH = "DI_COMPILER_TSCONFIG_PATH";
 export const ENV_VARIABLE_IDENTIFIER = "DI_COMPILER_IDENTIFIER";
 export const ENV_VARIABLE_DISABLE_CACHE = "DI_COMPILER_DISABLE_CACHE";
+export const ENV_VARIABLE_CLASS_NAME = "DI_COMPILER_CLASS_NAME";
+
 // Only these formats have type information that can be transpiled with DICompiler
 export const ALLOWED_EXTENSIONS = new Set([".ts", ".mts", ".cts"]);
 
@@ -18,6 +21,11 @@ interface DiTsconfigOptions {
 	 * Providing one or more identifiers up front can be considered an optimization, as this step can be skipped that way
 	 */
 	identifier?: MaybeArray<string>;
+	/*
+	 * Additional class name(s) that should be considered as the implementation class DIContainer. Required when inheriting
+	 * from DIContainer to add functionality to have the transformer recognize the class as a DIContainer.
+	 */
+	diClassName?: MaybeArray<string>;
 	disableCache: boolean;
 }
 
@@ -26,6 +34,8 @@ interface ExtendedTsconfig {
 	compilerOptions: TS.CompilerOptions;
 }
 
+let diClassName: MaybeArray<string> | undefined;
+
 export function resolveOptions(typescript: typeof TS): Partial<TransformOptions> {
 	const tsconfig = upgradeTsconfig(
 		typescript,
@@ -46,6 +56,15 @@ export function resolveOptions(typescript: typeof TS): Partial<TransformOptions>
 		identifier = identifier[0];
 	}
 
+	diClassName =
+		process.env[ENV_VARIABLE_CLASS_NAME]?.split(",")
+			.map(item => item.trim())
+			.filter(item => item.length > 0) ?? tsconfig.di?.diClassName;
+
+	if (Array.isArray(diClassName) && diClassName.length === 1) {
+		diClassName = diClassName[0];
+	}
+
 	const disableCache = process.env[ENV_VARIABLE_DISABLE_CACHE] == null ? tsconfig.di?.disableCache ?? false : booleanize(process.env[ENV_VARIABLE_DISABLE_CACHE]);
 
 	return {
@@ -57,6 +76,13 @@ export function resolveOptions(typescript: typeof TS): Partial<TransformOptions>
 	};
 }
 
+export function isDiClassName(name: string | TS.__String): boolean {
+	if (name === DI_CONTAINER_NAME) {
+		return true;
+	}
+	return diClassName == null ? false : Array.isArray(diClassName) ? diClassName.includes(<string> name) : diClassName === name;
+}
+
 function upgradeTsconfig(typescript: typeof TS, input?: TsConfigResult | TsConfigJsonResolved): ExtendedTsconfig {
 	if (input == null) {
 		return {

src/transformer/before/visitor/visit-call-expression.ts

@@ -1,11 +1,11 @@
-import {DI_CONTAINER_NAME} from "../../constant.js";
 import type {TS} from "../../../type/type.js";
 import type {BeforeVisitorOptions} from "../before-visitor-options.js";
 import type {DiMethodName} from "../../di-method-kind.js";
 import type {VisitorContext} from "../../visitor-context.js";
 import {getImportDefaultHelper, getImportStarHelper, moduleKindDefinesDependencies, moduleKindSupportsImportHelpers} from "../../../util/ts-util.js";
 import {pickServiceOrImplementationName} from "../util.js";
 import {ensureArray} from "../../../util/util.js";
+import { isDiClassName } from '../../../loader/shared.js';
 
 export function visitCallExpression(options: BeforeVisitorOptions<TS.CallExpression>): TS.VisitResult<TS.Node> {
 	const {node, childContinuation, continuation, context, addTslibDefinition, requireImportedSymbol} = options;
@@ -283,7 +283,7 @@ function isDiContainerInstance(node: TS.PropertyAccessExpression | TS.ElementAcc
 		// Don't proceed unless the left-hand expression is the DIServiceContainer
 		const type = context.typeChecker.getTypeAtLocation(node.expression);
 
-		if (type == null || type.symbol == null || type.symbol.escapedName !== DI_CONTAINER_NAME) {
+		if (type == null || type.symbol == null || !isDiClassName(type.symbol.escapedName)) {
 			return false;
 		}
 	} else {
@@ -303,7 +303,7 @@ function isDiContainerInstance(node: TS.PropertyAccessExpression | TS.ElementAcc
 				!evaluationResult.success ||
 				evaluationResult.value == null ||
 				typeof evaluationResult.value !== "object" ||
-				evaluationResult.value.constructor?.name !== DI_CONTAINER_NAME
+				!isDiClassName(evaluationResult.value.constructor?.name)
 			) {
 				return false;
 			}

src/transformer/di-options.ts

@@ -5,7 +5,7 @@ interface DiOptionsBase {
 }
 
 export interface DiProgramOptions extends DiOptionsBase {
-	program: TS.Program;
+	program: TS.Program & DiIsolatedModulesOptions;
 }
 
 export interface DiIsolatedModulesOptions extends DiOptionsBase {
@@ -15,6 +15,11 @@ export interface DiIsolatedModulesOptions extends DiOptionsBase {
 	 * Providing one or more identifiers up front can be considered an optimization, as this step can be skipped that way
 	 */
 	identifier?: MaybeArray<string>;
+	/*
+	 * Additional class name(s) that should be considered as the implementation class DIContainer. Required when inheriting
+	 * from DIContainer to add functionality to have the transofmrer recognize the class as a DIContainer.
+	 */
+	diClassName?: MaybeArray<string>;
 	compilerOptions?: TS.CompilerOptions;
 }
 

src/transformer/di.ts

@@ -1,13 +1,24 @@
-import type {TS} from "../type/type.js";
-import type {DiOptions} from "./di-options.js";
+import {TS} from "../type/type.js";
+import type {DiIsolatedModulesOptions, DiOptions} from "./di-options.js";
 import {beforeTransformer} from "./before/before-transformer.js";
 import {afterTransformer} from "./after/after-transformer.js";
 import {getBaseVisitorContext} from "./get-base-visitor-context.js";
+import { resolveOptions } from '../loader/shared.js';
 
 /**
  * CustomTransformer that associates constructor arguments with any given class declaration
  */
 export function di(options: DiOptions): TS.CustomTransformers {
+	const transformOptions = <DiIsolatedModulesOptions> resolveOptions(options.typescript ?? TS);
+
+	if ("program" in options) {
+		options.program.identifier = transformOptions.identifier;
+		options.program.diClassName = transformOptions.diClassName;
+	} else {
+		options.identifier = transformOptions.identifier;
+		options.diClassName = transformOptions.diClassName;
+	}
+
 	const baseVisitorContext = getBaseVisitorContext(options);
 
 	return {

src/transformer/get-base-visitor-context.ts

@@ -45,7 +45,8 @@ export function getBaseVisitorContext({typescript = TSModule as typeof TS, ...re
 	} else {
 		const compilerOptions = rest.compilerOptions ?? typescript.getDefaultCompilerOptions();
 		return {
-			identifier: [],
+			identifier: rest.identifier ?? [],
+			diClassName: rest.diClassName ?? [],
 
 			...rest,
 			...visitorContextShared,