Add Type Imports for Python

.chronus/changes/feature-python-type-imports-2026-2-16.md

@@ -0,0 +1,8 @@
+---
+changeKind: feature
+packages:
+  - "@alloy-js/python"
+---
+
+Add type-only imports support in Python code generation.
+

packages/python/src/components/CallSignature.tsx

@@ -13,6 +13,7 @@ import {
 import { createPythonSymbol } from "../symbol-creation.js";
 import { PythonOutputSymbol } from "../symbols/index.js";
 import { Atom } from "./Atom.jsx";
+import { TypeRefContext } from "./TypeRefContext.jsx";
 
 export type ParameterMarker = "*" | "/";
 
@@ -75,7 +76,10 @@ function parameter(param: DeclaredParameterDescriptor) {
     <group>
       {param.symbol.name}
       <Show when={!!param.type}>
-        : <TypeSlot>{param.type}</TypeSlot>
+        :{" "}
+        <TypeRefContext>
+          <TypeSlot>{param.type}</TypeSlot>
+        </TypeRefContext>
       </Show>
       <Show when={param.default !== undefined}>
         <Show when={!param.type}>=</Show>
@@ -184,7 +188,7 @@ export function CallSignature(props: CallSignatureProps) {
     props.returnType ?
       <>
         {" -> "}
-        {props.returnType}
+        <TypeRefContext>{props.returnType}</TypeRefContext>
       </>
     : undefined;
 

packages/python/src/components/ImportStatement.tsx

@@ -6,6 +6,53 @@ export interface ImportStatementsProps {
   joinImportsFromSameModule?: boolean;
 }
 
+export interface CategorizedImports {
+  /** Imports used only in type annotation contexts (for TYPE_CHECKING block) */
+  typeImports: ImportRecords;
+  /** Regular imports used at runtime */
+  valueImports: ImportRecords;
+}
+
+/**
+ * Categorize import records into type-only and value imports.
+ * Type-only imports are those used only in type annotation contexts.
+ * Value imports are regular imports used at runtime.
+ */
+export function categorizeImportRecords(
+  records: ImportRecords,
+): CategorizedImports {
+  const typeImports = new Map() as ImportRecords;
+  const valueImports = new Map() as ImportRecords;
+
+  for (const [module, properties] of records) {
+    if (!properties.symbols || properties.symbols.size === 0) {
+      // Module-level imports without symbols go to value imports
+      valueImports.set(module, properties);
+      continue;
+    }
+
+    const typeSymbols = new Set<ImportedSymbol>();
+    const valueSymbols = new Set<ImportedSymbol>();
+
+    for (const sym of properties.symbols) {
+      if (sym.local.isTypeOnly) {
+        typeSymbols.add(sym);
+      } else {
+        valueSymbols.add(sym);
+      }
+    }
+
+    if (typeSymbols.size > 0) {
+      typeImports.set(module, { symbols: typeSymbols });
+    }
+    if (valueSymbols.size > 0) {
+      valueImports.set(module, { symbols: valueSymbols });
+    }
+  }
+
+  return { typeImports, valueImports };
+}
+
 /**
  * A component that renders import statements based on the provided import records.
  *
@@ -16,11 +63,11 @@ export interface ImportStatementsProps {
  */
 export function ImportStatements(props: ImportStatementsProps) {
   // Sort the import records by module name
-  const imports = computed(() =>
-    [...props.records].sort(([a], [b]) => {
+  const imports = computed(() => {
+    return [...props.records].sort(([a], [b]) => {
       return a.name.localeCompare(b.name);
-    }),
-  );
+    });
+  });
 
   return mapJoin(
     () => imports.value,

packages/python/src/components/Reference.tsx

@@ -1,4 +1,5 @@
 import { computed, emitSymbol, Refkey } from "@alloy-js/core";
+import { isTypeRefContext } from "../context/type-ref-context.js";
 import { ref } from "../symbols/index.js";
 
 export interface ReferenceProps {
@@ -13,7 +14,8 @@ export interface ReferenceProps {
  * It takes a `refkey` prop which is the key of the symbol to reference.
  */
 export function Reference({ refkey }: ReferenceProps) {
-  const reference = ref(refkey);
+  const inTypeRef = isTypeRefContext();
+  const reference = ref(refkey, { type: inTypeRef });
   const symbolRef = computed(() => reference()[1]);
 
   emitSymbol(symbolRef);

packages/python/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

packages/python/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>
+  );
+};

packages/python/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>
   );
 }

packages/python/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>
       </>
     );
   });

packages/python/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";

packages/python/src/context/index.ts

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

packages/python/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;
+}