diff --git a/.chronus/changes/feature-python-type-imports-2026-2-16.md b/.chronus/changes/feature-python-type-imports-2026-2-16.md new file mode 100644 index 00000000..4d383d94 --- /dev/null +++ b/.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. + diff --git a/packages/python/src/components/CallSignature.tsx b/packages/python/src/components/CallSignature.tsx index c09ad125..e19214e8 100644 --- a/packages/python/src/components/CallSignature.tsx +++ b/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) { {param.symbol.name} - : {param.type} + :{" "} + + {param.type} + = @@ -184,7 +188,7 @@ export function CallSignature(props: CallSignatureProps) { props.returnType ? <> {" -> "} - {props.returnType} + {props.returnType} : undefined; diff --git a/packages/python/src/components/ImportStatement.tsx b/packages/python/src/components/ImportStatement.tsx index 35ca1e93..09d54b24 100644 --- a/packages/python/src/components/ImportStatement.tsx +++ b/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(); + const valueSymbols = new Set(); + + 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, diff --git a/packages/python/src/components/Reference.tsx b/packages/python/src/components/Reference.tsx index 6ac0a562..eadd2ee0 100644 --- a/packages/python/src/components/Reference.tsx +++ b/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); diff --git a/packages/python/src/components/SourceFile.tsx b/packages/python/src/components/SourceFile.tsx index 3dd517d8..0c7999cd 100644 --- a/packages/python/src/components/SourceFile.tsx +++ b/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 ( - 0}> - - - - - + {/* Regular (non-type-only) imports */} + 0}> + + + + {/* TYPE_CHECKING block with type-only imports */} + 0}> + + + + + + {/* Spacing after imports */} + 0}> + {/* Extra blank line before top-level definitions */} + * {someTypeRefkey} + * + * ``` + */ +export const TypeRefContext = ({ children }: TypeRefContextProps) => { + return ( + {children} + ); +}; diff --git a/packages/python/src/components/TypeReference.tsx b/packages/python/src/components/TypeReference.tsx index e0c014a4..8b61b741 100644 --- a/packages/python/src/components/TypeReference.tsx +++ b/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 ( - - + + + + + {type} + {typeArgs} + - {type} - {typeArgs} - - - + + ); } diff --git a/packages/python/src/components/VariableDeclaration.tsx b/packages/python/src/components/VariableDeclaration.tsx index 0e352ae3..c68234a6 100644 --- a/packages/python/src/components/VariableDeclaration.tsx +++ b/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 ( <> - : {props.type} + :{" "} + + {props.type} + ); }); diff --git a/packages/python/src/components/index.ts b/packages/python/src/components/index.ts index 32f69840..710a1920 100644 --- a/packages/python/src/components/index.ts +++ b/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"; diff --git a/packages/python/src/context/index.ts b/packages/python/src/context/index.ts new file mode 100644 index 00000000..36b1ed77 --- /dev/null +++ b/packages/python/src/context/index.ts @@ -0,0 +1 @@ +export * from "./type-ref-context.js"; diff --git a/packages/python/src/context/type-ref-context.tsx b/packages/python/src/context/type-ref-context.tsx new file mode 100644 index 00000000..39444944 --- /dev/null +++ b/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 = createContext(); + +/** + * @returns 'true' if in a type context, 'false' if in a value context. + */ +export function isTypeRefContext(): boolean { + return useContext(TypeRefContextDef) === true; +} diff --git a/packages/python/src/symbols/python-module-scope.ts b/packages/python/src/symbols/python-module-scope.ts index fea16b3a..870ea9c2 100644 --- a/packages/python/src/symbols/python-module-scope.ts +++ b/packages/python/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 {} +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 = 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 { return { ...super.debugInfo, diff --git a/packages/python/src/symbols/python-output-symbol.ts b/packages/python/src/symbols/python-output-symbol.ts index 7678b4f1..4c7a6f69 100644 --- a/packages/python/src/symbols/python-output-symbol.ts +++ b/packages/python/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); diff --git a/packages/python/src/symbols/reference.tsx b/packages/python/src/symbols/reference.tsx index c5a99ced..ca2182c8 100644 --- a/packages/python/src/symbols/reference.tsx +++ b/packages/python/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( @@ -41,6 +50,7 @@ export function ref( sourceFile!.scope.addImport( lexicalDeclaration, pathDown[0] as PythonModuleScope, + { type: options?.type }, ), ); } diff --git a/packages/python/test/dataclassdeclarations.test.tsx b/packages/python/test/dataclassdeclarations.test.tsx index 03fbb23f..d698e3b1 100644 --- a/packages/python/test/dataclassdeclarations.test.tsx +++ b/packages/python/test/dataclassdeclarations.test.tsx @@ -73,7 +73,10 @@ describe("DataclassDeclaration", () => { expect(res).toRenderTo( d` from dataclasses import dataclass - from dataclasses import KW_ONLY + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from dataclasses import KW_ONLY @dataclass @@ -624,7 +627,10 @@ describe("DataclassDeclaration", () => { name="user" type={userRefkey} initializer={ - + } /> diff --git a/packages/python/test/externals.test.tsx b/packages/python/test/externals.test.tsx index b59af614..1f8f285f 100644 --- a/packages/python/test/externals.test.tsx +++ b/packages/python/test/externals.test.tsx @@ -92,7 +92,10 @@ it("uses import from external library in multiple functions", () => { const expected = d` from requests import get from requests import post - from requests.models import Response + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from requests.models import Response def get_user(user_id: int) -> Response: @@ -169,7 +172,10 @@ it("uses import from external library in multiple class methods", () => { const expected = d` from requests import get from requests import post - from requests.models import Response + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from requests.models import Response class UserClient: diff --git a/packages/python/test/functiondeclaration.test.tsx b/packages/python/test/functiondeclaration.test.tsx index 54e59664..de06d69f 100644 --- a/packages/python/test/functiondeclaration.test.tsx +++ b/packages/python/test/functiondeclaration.test.tsx @@ -450,9 +450,12 @@ describe("Function Declaration", () => { `, "usage.py": ` - from mod1 import Foo - from mod2 import A - from mod2 import B + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from mod1 import Foo + from mod2 import A + from mod2 import B async def foo(x: A, y: B, *args, **kwargs) -> Foo: diff --git a/packages/python/test/imports.test.tsx b/packages/python/test/imports.test.tsx index c2e8e23e..1304e368 100644 --- a/packages/python/test/imports.test.tsx +++ b/packages/python/test/imports.test.tsx @@ -1,6 +1,9 @@ import { refkey } from "@alloy-js/core"; import { describe, expect, it } from "vitest"; -import { ImportStatement } from "../src/components/ImportStatement.jsx"; +import { + ImportStatement, + ImportStatements, +} from "../src/components/ImportStatement.jsx"; import * as py from "../src/index.js"; import { createPythonSymbol } from "../src/symbol-creation.js"; import { ImportedSymbol, ImportRecords } from "../src/symbols/index.js"; @@ -56,7 +59,7 @@ describe("ImportStatements", () => { [sysModuleScope, { symbols: new Set() }], ]); - const result = toSourceText([]); + const result = toSourceText([]); const expected = ` from math import pi from math import sqrt @@ -85,10 +88,7 @@ describe("ImportStatements", () => { ]); const result = toSourceText([ - , + , ]); const expected = ` from math import pi, sqrt diff --git a/packages/python/test/references.test.tsx b/packages/python/test/references.test.tsx index 7e3e1560..5a2abde1 100644 --- a/packages/python/test/references.test.tsx +++ b/packages/python/test/references.test.tsx @@ -15,7 +15,7 @@ describe("Reference", () => { name="current_user" type={rk1} initializer={ - + } /> , diff --git a/packages/python/test/sourcefiles.test.tsx b/packages/python/test/sourcefiles.test.tsx index 4b752111..1bd80ab4 100644 --- a/packages/python/test/sourcefiles.test.tsx +++ b/packages/python/test/sourcefiles.test.tsx @@ -312,7 +312,7 @@ it("only futureImports before definition", () => { const content = ( } + futureImports={[]} > pass @@ -333,7 +333,7 @@ it("only futureImports before non-definition", () => { const content = ( } + futureImports={[]} > @@ -393,7 +393,7 @@ it("doc + futureImports before definition", () => { } + futureImports={[]} > pass @@ -423,7 +423,7 @@ it("doc + futureImports before non-definition", () => { } + futureImports={[]} > @@ -498,7 +498,7 @@ it("futureImports + imports before definition", () => { const content = ( } + futureImports={[]} > @@ -525,7 +525,7 @@ it("futureImports + imports before non-definition", () => { const content = ( } + futureImports={[]} > { } + futureImports={[]} > @@ -588,7 +588,7 @@ it("doc + futureImports + imports before non-definition", () => { } + futureImports={[]} > { const content = ( } + futureImports={[]} /> ); @@ -661,7 +661,7 @@ it("doc + futureImports in file (no children)", () => { } + futureImports={[]} /> ); @@ -743,7 +743,7 @@ it("headerComment + futureImports before definition", () => { } + futureImports={[]} > pass @@ -797,7 +797,7 @@ it("headerComment + doc + futureImports before definition", () => { path="test.py" headerComment="Copyright 2024 My Company" doc={moduleDoc} - futureImports={} + futureImports={[]} > pass @@ -829,7 +829,7 @@ it("headerComment + doc + futureImports + imports before definition", () => { path="test.py" headerComment="Copyright 2024 My Company" doc={moduleDoc} - futureImports={} + futureImports={[]} > diff --git a/packages/python/test/type-checking-imports.test.tsx b/packages/python/test/type-checking-imports.test.tsx new file mode 100644 index 00000000..8b743fa6 --- /dev/null +++ b/packages/python/test/type-checking-imports.test.tsx @@ -0,0 +1,363 @@ +import { refkey } from "@alloy-js/core"; +import { describe, it } from "vitest"; +import { createModule } from "../src/create-module.js"; +import * as py from "../src/index.js"; +import { assertFileContents, toSourceTextMultiple } from "./utils.jsx"; + +describe("TYPE_CHECKING imports", () => { + it("imports type-only references inside TYPE_CHECKING block", () => { + const userClassRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + pass + + , + ]); + + assertFileContents(result, { + "models.py": ` + class User: + pass + + `, + "service.py": ` + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from models import User + + + def process_user(user: User) -> None: + pass + + `, + }); + }); + + it("imports value references outside TYPE_CHECKING block", () => { + const userClassRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + } + /> + + , + ]); + + assertFileContents(result, { + "models.py": ` + class User: + pass + + `, + "service.py": ` + from models import User + + + def create_user() -> None: + user = User() + + `, + }); + }); + + it("upgrades type-only import to regular import when also used as value", () => { + const userClassRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + + } + /> + <>return user + + + , + ]); + + // Since User is used both as a type (parameter type, return type) and + // as a value (ClassInstantiation), it should be a regular import + assertFileContents(result, { + "models.py": ` + class User: + pass + + `, + "service.py": ` + from models import User + + + def create_user(existing: User) -> User: + user = User() + return user + + `, + }); + }); + + it("handles mixed type-only and regular imports from same module", () => { + const userClassRef = refkey(); + const helperFuncRef = refkey(); + + const result = toSourceTextMultiple([ + + + + pass + + , + + + + + , + ]); + + // helper is used as a value (function call), User is only used as type + assertFileContents(result, { + "models.py": ` + class User: + pass + + + def helper(): + pass + + `, + "service.py": ` + from models import helper + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from models import User + + + def process(user: User) -> None: + helper() + + `, + }); + }); + + it("handles return type as type-only import", () => { + const resultTypeRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + pass + + , + ]); + + assertFileContents(result, { + "types.py": ` + class Result: + pass + + `, + "main.py": ` + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from types import Result + + + def get_result() -> Result: + pass + + `, + }); + }); + + it("handles variable type annotation as type-only import", () => { + const configTypeRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + , + ]); + + assertFileContents(result, { + "types.py": ` + class Config: + pass + + `, + "main.py": ` + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from types import Config + + config: Config + `, + }); + }); + + it("handles TypeReference component as type-only import", () => { + const myTypeRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + } + omitNone + /> + , + ]); + + assertFileContents(result, { + "types.py": ` + class MyType: + pass + + `, + "main.py": ` + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from types import MyType + + value: MyType + `, + }); + }); + + it("handles class bases as regular import (runtime requirement)", () => { + const baseClassRef = refkey(); + + const result = toSourceTextMultiple([ + + + , + + + , + ]); + + // Class bases require runtime access, so they should NOT be + // inside a TYPE_CHECKING block + assertFileContents(result, { + "base.py": ` + class BaseClass: + pass + + `, + "derived.py": ` + from base import BaseClass + + + class DerivedClass(BaseClass): + pass + + `, + }); + }); + + it("renders regular imports before TYPE_CHECKING block", () => { + // Create a typing module with multiple exports + const typingModule = createModule({ + name: "typing", + descriptor: { + ".": ["TYPE_CHECKING", "cast"], + }, + }); + + // Create a third-party module + const requestsModule = createModule({ + name: "requests", + descriptor: { + ".": ["get"], + }, + }); + + const userClassRef = refkey(); + + const result = toSourceTextMultiple( + [ + + + , + + + + {/* Use cast as a value (function call) to make it a regular import */} + <>response = {requestsModule["."].get}("https://example.com") + <>return {typingModule["."].cast}(str, user) + + + , + ], + { externals: [typingModule, requestsModule] }, + ); + + // Regular imports first (sorted alphabetically), then TYPE_CHECKING block + assertFileContents(result, { + "models.py": ` + class User: + pass + + `, + "service.py": ` + from requests import get + from typing import cast + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from models import User + + + def get_users(user: User) -> str: + response = get("https://example.com") + return cast(str, user) + + `, + }); + }); +}); diff --git a/packages/python/test/uniontypeexpression.test.tsx b/packages/python/test/uniontypeexpression.test.tsx index ca75fd86..3a23f861 100644 --- a/packages/python/test/uniontypeexpression.test.tsx +++ b/packages/python/test/uniontypeexpression.test.tsx @@ -139,7 +139,10 @@ describe("UnionTypeExpression", () => { `, "use.py": ` - from defs import Bar + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from defs import Bar v: Bar[T] = None `, diff --git a/packages/python/test/variables.test.tsx b/packages/python/test/variables.test.tsx index 89b94bc6..6ad9517d 100644 --- a/packages/python/test/variables.test.tsx +++ b/packages/python/test/variables.test.tsx @@ -183,7 +183,10 @@ describe("Python Variable", () => { `, "usage.py": ` - from classes import MyClass + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from classes import MyClass my_var: MyClass = None `, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d04fe6d0..4027d2fb 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -447,12 +447,6 @@ importers: specifier: 'catalog:' version: 3.2.4(@types/debug@4.1.12)(@types/node@24.10.9)(esbuild@0.25.8)(jiti@2.6.1)(tsx@4.20.3)(yaml@2.8.0) - packages/dev-tools: - dependencies: - '@radix-ui/react-context-menu': - specifier: ^2.2.16 - version: 2.2.16(@types/react-dom@19.2.3(@types/react@19.2.8))(@types/react@19.2.8)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) - packages/devtools: dependencies: '@radix-ui/react-context-menu':