@alloy-js/python 0.4.0-dev.3 → 0.4.0-dev.6

package/src/symbols/python-module-scope.ts

@@ -2,6 +2,30 @@ import { createSymbol, reactive, shallowReactive } from "@alloy-js/core";
 import { PythonLexicalScope } from "./python-lexical-scope.js";
 import { PythonOutputSymbol } from "./python-output-symbol.js";
 
+// Internal typing module for TYPE_CHECKING imports
+let _typingModuleScope: PythonModuleScope | undefined;
+let _typeCheckingSymbol: PythonOutputSymbol | undefined;
+
+/**
+ * Get the internal typing module scope and TYPE_CHECKING symbol.
+ * Used by addTypeImport() to add TYPE_CHECKING imports without
+ * going through the binder's refkey resolution.
+ */
+function getTypingModuleInternal(): {
+  scope: PythonModuleScope;
+  TYPE_CHECKING: PythonOutputSymbol;
+} {
+  if (!_typingModuleScope) {
+    _typingModuleScope = new PythonModuleScope("typing", undefined);
+    _typeCheckingSymbol = new PythonOutputSymbol(
+      "TYPE_CHECKING",
+      _typingModuleScope.symbols,
+      {},
+    );
+  }
+  return { scope: _typingModuleScope, TYPE_CHECKING: _typeCheckingSymbol! };
+}
+
 export class ImportedSymbol {
   local: PythonOutputSymbol;
   target: PythonOutputSymbol;
@@ -22,6 +46,14 @@ export interface ImportRecordProps {
 
 export class ImportRecords extends Map<PythonModuleScope, ImportRecordProps> {}
 
+export interface AddImportOptions {
+  /**
+   * If true, this import is only used in type annotation contexts.
+   * Such imports will be guarded with `if TYPE_CHECKING:`.
+   */
+  type?: boolean;
+}
+
 export class PythonModuleScope extends PythonLexicalScope {
   #importedSymbols: Map<PythonOutputSymbol, PythonOutputSymbol> =
     shallowReactive(new Map());
@@ -34,9 +66,17 @@ export class PythonModuleScope extends PythonLexicalScope {
     return this.#importedModules;
   }
 
-  addImport(targetSymbol: PythonOutputSymbol, targetModule: PythonModuleScope) {
+  addImport(
+    targetSymbol: PythonOutputSymbol,
+    targetModule: PythonModuleScope,
+    options?: AddImportOptions,
+  ) {
     const existing = this.importedSymbols.get(targetSymbol);
     if (existing) {
+      // If existing is type-only but now used as value, upgrade it
+      if (!options?.type && existing.isTypeOnly) {
+        existing.markAsValue();
+      }
       return existing;
     }
 
@@ -50,7 +90,11 @@ export class PythonModuleScope extends PythonLexicalScope {
       PythonOutputSymbol,
       targetSymbol.name,
       this.symbols,
-      { binder: this.binder, aliasTarget: targetSymbol },
+      {
+        binder: this.binder,
+        aliasTarget: targetSymbol,
+        typeOnly: options?.type,
+      },
     );
 
     this.importedSymbols.set(targetSymbol, localSymbol);
@@ -62,6 +106,15 @@ export class PythonModuleScope extends PythonLexicalScope {
     return localSymbol;
   }
 
+  /**
+   * Add TYPE_CHECKING import from the typing module.
+   * Returns the local symbol for use in the if block opener.
+   */
+  addTypeImport(): PythonOutputSymbol {
+    const typing = getTypingModuleInternal();
+    return this.addImport(typing.TYPE_CHECKING, typing.scope);
+  }
+
   override get debugInfo(): Record<string, unknown> {
     return {
       ...super.debugInfo,

package/src/symbols/python-output-symbol.ts

@@ -1,13 +1,19 @@
 import {
+  createSymbol,
   Namekey,
   OutputSpace,
   OutputSymbol,
   OutputSymbolOptions,
-  createSymbol,
+  track,
+  TrackOpTypes,
+  trigger,
+  TriggerOpTypes,
 } from "@alloy-js/core";
 
 export interface PythonOutputSymbolOptions extends OutputSymbolOptions {
   module?: string;
+  /** Whether this symbol is only used in type annotation contexts */
+  typeOnly?: boolean;
 }
 
 export interface CreatePythonSymbolFunctionOptions
@@ -28,6 +34,7 @@ export class PythonOutputSymbol extends OutputSymbol {
   ) {
     super(name, spaces, options);
     this.#module = options.module ?? undefined;
+    this.#typeOnly = options.typeOnly ?? false;
   }
 
   // The module in which the symbol is defined
@@ -37,6 +44,29 @@ export class PythonOutputSymbol extends OutputSymbol {
     return this.#module;
   }
 
+  #typeOnly: boolean;
+
+  /**
+   * Returns true if this symbol is only used in type annotation contexts.
+   * Such symbols can be imported inside a TYPE_CHECKING block.
+   */
+  get isTypeOnly() {
+    track(this, TrackOpTypes.GET, "typeOnly");
+    return this.#typeOnly;
+  }
+
+  /**
+   * Mark this symbol as also being used as a value (not just a type).
+   */
+  markAsValue() {
+    if (!this.#typeOnly) {
+      return;
+    }
+    const oldValue = this.#typeOnly;
+    this.#typeOnly = false;
+    trigger(this, TriggerOpTypes.SET, "typeOnly", false, oldValue);
+  }
+
   get staticMembers() {
     return this.memberSpaceFor("static")!;
   }
@@ -69,6 +99,7 @@ export class PythonOutputSymbol extends OutputSymbol {
       aliasTarget: this.aliasTarget,
       module: this.module,
       metadata: this.metadata,
+      typeOnly: this.isTypeOnly,
     });
 
     this.initializeCopy(copy);

package/src/symbols/reference.tsx

@@ -13,8 +13,17 @@ import { PythonModuleScope } from "./python-module-scope.js";
 import { PythonOutputSymbol } from "./python-output-symbol.js";
 import { PythonOutputScope } from "./scopes.js";
 
+export interface RefOptions {
+  /**
+   * If true, this reference is only used in a type annotation context.
+   * The import will be guarded with `if TYPE_CHECKING:`.
+   */
+  type?: boolean;
+}
+
 export function ref(
   refkey: Refkey,
+  options?: RefOptions,
 ): () => [Children, PythonOutputSymbol | undefined] {
   const sourceFile = useContext(PythonSourceFileContext);
   const resolveResult = resolve<PythonOutputScope, PythonOutputSymbol>(
@@ -41,6 +50,7 @@ export function ref(
         sourceFile!.scope.addImport(
           lexicalDeclaration,
           pathDown[0] as PythonModuleScope,
+          { type: options?.type },
         ),
       );
     }