@alloy-js/python 0.4.0-dev.2 → 0.4.0-dev.5

package/src/components/SourceFile.tsx

@@ -1,6 +1,7 @@
 import {
   childrenArray,
   ComponentContext,
+  computed,
   SourceFile as CoreSourceFile,
   createNamedContext,
   createScope,
@@ -14,8 +15,12 @@ import {
 } from "@alloy-js/core";
 import { join } from "pathe";
 import { PythonModuleScope } from "../symbols/index.js";
-import { ImportStatements } from "./ImportStatement.js";
+import {
+  categorizeImportRecords,
+  ImportStatements,
+} from "./ImportStatement.js";
 import { SimpleCommentBlock } from "./PyDoc.js";
+import { PythonBlock } from "./PythonBlock.js";
 import { Reference } from "./Reference.js";
 
 // Non top-level definitions
@@ -99,7 +104,7 @@ export interface SourceFileProps {
   /**
    * __future__ imports to render after the docstring but before regular imports.
    */
-  futureImports?: Children;
+  futureImports?: Children[];
 }
 
 /**
@@ -167,6 +172,28 @@ export function SourceFile(props: SourceFileProps) {
     props.doc !== undefined ||
     props.futureImports !== undefined;
 
+  const imports = computed(() => {
+    // Quick scan for any type-only imports
+    const hasTypeImports = [...scope.importedModules.values()].some(
+      (props) =>
+        props.symbols && [...props.symbols].some((s) => s.local.isTypeOnly),
+    );
+
+    // Add TYPE_CHECKING before categorizing so it's naturally included
+    const typeImportSymbol = hasTypeImports ? scope.addTypeImport() : undefined;
+
+    // Single categorize - TYPE_CHECKING is already in scope.importedModules
+    const { valueImports, typeImports } = categorizeImportRecords(
+      scope.importedModules,
+    );
+
+    return {
+      valueImports,
+      typeImports,
+      typeImportSymbol,
+    };
+  });
+
   return (
     <CoreSourceFile
       path={props.path}
@@ -216,12 +243,21 @@ export function SourceFile(props: SourceFileProps) {
           <hbr />
         </Show>
       </Show>
-      <Show when={scope.importedModules.size > 0}>
-        <ImportStatements records={scope.importedModules} />
-        <Show when={hasChildren}>
-          <hbr />
-          <hbr />
-        </Show>
+      {/* Regular (non-type-only) imports */}
+      <Show when={imports.value.valueImports.size > 0}>
+        <ImportStatements records={imports.value.valueImports} />
+        <hbr />
+      </Show>
+      {/* TYPE_CHECKING block with type-only imports */}
+      <Show when={imports.value.typeImports.size > 0}>
+        <hbr />
+        <PythonBlock opener={`if ${imports.value.typeImportSymbol!.name}:`}>
+          <ImportStatements records={imports.value.typeImports} />
+        </PythonBlock>
+      </Show>
+      {/* Spacing after imports */}
+      <Show when={hasChildren && scope.importedModules.size > 0}>
+        <hbr />
       </Show>
       {/* Extra blank line before top-level definitions */}
       <Show

package/src/components/TypeRefContext.tsx

@@ -0,0 +1,36 @@
+import type { Children } from "@alloy-js/core";
+import {
+  isTypeRefContext,
+  TypeRefContextDef,
+} from "../context/type-ref-context.js";
+
+// Re-export for external use
+export { isTypeRefContext };
+
+export interface TypeRefContextProps {
+  /**
+   * Children
+   */
+  children: Children;
+}
+
+/**
+ * Set the current context of reference to be type reference.
+ *
+ * @remarks
+ * References used inside the children of this component will be treated as
+ * type-only references. When a symbol is only referenced in type contexts,
+ * it will be imported inside a `if TYPE_CHECKING:` block.
+ *
+ * @example
+ * ```tsx
+ * <TypeRefContext>
+ *   {someTypeRefkey}
+ * </TypeRefContext>
+ * ```
+ */
+export const TypeRefContext = ({ children }: TypeRefContextProps) => {
+  return (
+    <TypeRefContextDef.Provider value>{children}</TypeRefContextDef.Provider>
+  );
+};

package/src/components/TypeReference.tsx

@@ -1,5 +1,6 @@
 import { Children, Refkey, Show } from "@alloy-js/core";
 import { TypeArguments } from "./TypeArguments.js";
+import { TypeRefContext } from "./TypeRefContext.js";
 
 export interface TypeReferenceProps {
   /** A refkey to a declared symbol. */
@@ -12,6 +13,11 @@ export interface TypeReferenceProps {
 
 /**
  * A type reference like Foo[T, P] or int.
+ *
+ * @remarks
+ * This component automatically wraps its content in a type reference context,
+ * so any symbols referenced via refkey will be imported as type-only
+ * (inside a `if TYPE_CHECKING:` block) unless also used as values elsewhere.
  */
 export function TypeReference(props: TypeReferenceProps) {
   const type = props.refkey ? props.refkey : props.name;
@@ -21,13 +27,15 @@ export function TypeReference(props: TypeReferenceProps) {
     : undefined;
 
   return (
-    <group>
-      <indent>
+    <TypeRefContext>
+      <group>
+        <indent>
+          <sbr />
+          {type}
+          <Show when={Boolean(typeArgs)}>{typeArgs}</Show>
+        </indent>
         <sbr />
-        {type}
-        <Show when={Boolean(typeArgs)}>{typeArgs}</Show>
-      </indent>
-      <sbr />
-    </group>
+      </group>
+    </TypeRefContext>
   );
 }

package/src/components/VariableDeclaration.tsx

@@ -10,6 +10,7 @@ import { createPythonSymbol } from "../symbol-creation.js";
 import { Atom } from "./Atom.jsx";
 import { BaseDeclarationProps } from "./Declaration.jsx";
 import { SimpleCommentBlock } from "./PyDoc.jsx";
+import { TypeRefContext } from "./TypeRefContext.jsx";
 
 export interface VariableDeclarationProps extends BaseDeclarationProps {
   /**
@@ -93,7 +94,10 @@ export function VariableDeclaration(props: VariableDeclarationProps) {
     if (!props.type || props.callStatementVar) return undefined;
     return (
       <>
-        : <TypeSymbolSlot>{props.type}</TypeSymbolSlot>
+        :{" "}
+        <TypeRefContext>
+          <TypeSymbolSlot>{props.type}</TypeSymbolSlot>
+        </TypeRefContext>
       </>
     );
   });

package/src/components/index.ts

@@ -13,7 +13,6 @@ export type { CommonFunctionProps } from "./FunctionBase.js";
 export * from "./FunctionCallExpression.js";
 export * from "./FunctionDeclaration.js";
 export * from "./FutureStatement.js";
-export * from "./ImportStatement.js";
 export * from "./LexicalScope.js";
 export * from "./MemberExpression.js";
 export * from "./MemberScope.jsx";
@@ -27,6 +26,7 @@ export * from "./SourceFile.js";
 export * from "./StatementList.js";
 export * from "./StaticMethodDeclaration.js";
 export * from "./TypeArguments.js";
+export * from "./TypeRefContext.js";
 export * from "./TypeReference.js";
 export * from "./UnionTypeExpression.js";
 export * from "./VariableDeclaration.js";

package/src/context/index.ts

@@ -0,0 +1 @@
+export * from "./type-ref-context.js";

package/src/context/type-ref-context.tsx

@@ -0,0 +1,16 @@
+import { ComponentContext, createContext, useContext } from "@alloy-js/core";
+
+/**
+ * Context for tracking whether we are in a type annotation position.
+ * Used to determine if imports should be guarded with TYPE_CHECKING.
+ *
+ * @internal Use the TypeRefContext component instead.
+ */
+export const TypeRefContextDef: ComponentContext<true> = createContext<true>();
+
+/**
+ * @returns 'true' if in a type context, 'false' if in a value context.
+ */
+export function isTypeRefContext(): boolean {
+  return useContext(TypeRefContextDef) === true;
+}

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 },
         ),
       );
     }