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

package/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={
-                <py.ClassInstantiation target="User" args={["1", '"Alice"']} />
+                <py.ClassInstantiation
+                  target={userRefkey}
+                  args={["1", '"Alice"']}
+                />
               }
             />
             <hbr />

package/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:

package/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:

package/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<ImportedSymbol>() }],
     ]);
 
-    const result = toSourceText([<py.ImportStatements records={records} />]);
+    const result = toSourceText([<ImportStatements records={records} />]);
     const expected = `
     from math import pi
     from math import sqrt
@@ -85,10 +88,7 @@ describe("ImportStatements", () => {
     ]);
 
     const result = toSourceText([
-      <py.ImportStatements
-        records={records}
-        joinImportsFromSameModule={true}
-      />,
+      <ImportStatements records={records} joinImportsFromSameModule={true} />,
     ]);
     const expected = `
     from math import pi, sqrt

package/test/references.test.tsx

@@ -15,7 +15,7 @@ describe("Reference", () => {
           name="current_user"
           type={rk1}
           initializer={
-            <py.ClassInstantiation target="User" args={['"Marvin"']} />
+            <py.ClassInstantiation target={rk1} args={['"Marvin"']} />
           }
         />
       </py.SourceFile>,

package/test/sourcefiles.test.tsx

@@ -312,7 +312,7 @@ it("only futureImports before definition", () => {
   const content = (
     <py.SourceFile
       path="test.py"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.FunctionDeclaration name="hello">pass</py.FunctionDeclaration>
     </py.SourceFile>
@@ -333,7 +333,7 @@ it("only futureImports before non-definition", () => {
   const content = (
     <py.SourceFile
       path="test.py"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.VariableDeclaration name="x" initializer={42} />
     </py.SourceFile>
@@ -393,7 +393,7 @@ it("doc + futureImports before definition", () => {
     <py.SourceFile
       path="test.py"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.FunctionDeclaration name="hello">pass</py.FunctionDeclaration>
     </py.SourceFile>
@@ -423,7 +423,7 @@ it("doc + futureImports before non-definition", () => {
     <py.SourceFile
       path="test.py"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.VariableDeclaration name="x" initializer={42} />
     </py.SourceFile>
@@ -498,7 +498,7 @@ it("futureImports + imports before definition", () => {
   const content = (
     <py.SourceFile
       path="test.py"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.DataclassDeclaration name="User">
         <py.VariableDeclaration name="name" type="str" />
@@ -525,7 +525,7 @@ it("futureImports + imports before non-definition", () => {
   const content = (
     <py.SourceFile
       path="test.py"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.VariableDeclaration
         name="x"
@@ -552,7 +552,7 @@ it("doc + futureImports + imports before definition", () => {
     <py.SourceFile
       path="test.py"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.DataclassDeclaration name="User">
         <py.VariableDeclaration name="name" type="str" />
@@ -588,7 +588,7 @@ it("doc + futureImports + imports before non-definition", () => {
     <py.SourceFile
       path="test.py"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.VariableDeclaration
         name="x"
@@ -642,7 +642,7 @@ it("only futureImports in file (no children)", () => {
   const content = (
     <py.SourceFile
       path="test.py"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     />
   );
 
@@ -661,7 +661,7 @@ it("doc + futureImports in file (no children)", () => {
     <py.SourceFile
       path="test.py"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     />
   );
 
@@ -743,7 +743,7 @@ it("headerComment + futureImports before definition", () => {
     <py.SourceFile
       path="test.py"
       headerComment="Copyright 2024 My Company"
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.FunctionDeclaration name="hello">pass</py.FunctionDeclaration>
     </py.SourceFile>
@@ -797,7 +797,7 @@ it("headerComment + doc + futureImports before definition", () => {
       path="test.py"
       headerComment="Copyright 2024 My Company"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.FunctionDeclaration name="hello">pass</py.FunctionDeclaration>
     </py.SourceFile>
@@ -829,7 +829,7 @@ it("headerComment + doc + futureImports + imports before definition", () => {
       path="test.py"
       headerComment="Copyright 2024 My Company"
       doc={moduleDoc}
-      futureImports={<py.FutureStatement feature="annotations" />}
+      futureImports={[<py.FutureStatement feature="annotations" />]}
     >
       <py.DataclassDeclaration name="User">
         <py.VariableDeclaration name="name" type="str" />

package/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([
+      <py.SourceFile path="models.py">
+        <py.ClassDeclaration name="User" refkey={userClassRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="service.py">
+        <py.FunctionDeclaration
+          name="process_user"
+          parameters={[{ name: "user", type: userClassRef }]}
+          returnType="None"
+        >
+          pass
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+    ]);
+
+    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([
+      <py.SourceFile path="models.py">
+        <py.ClassDeclaration name="User" refkey={userClassRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="service.py">
+        <py.FunctionDeclaration name="create_user" returnType="None">
+          <py.VariableDeclaration
+            name="user"
+            initializer={<py.ClassInstantiation target={userClassRef} />}
+          />
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+    ]);
+
+    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([
+      <py.SourceFile path="models.py">
+        <py.ClassDeclaration name="User" refkey={userClassRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="service.py">
+        <py.FunctionDeclaration
+          name="create_user"
+          parameters={[{ name: "existing", type: userClassRef }]}
+          returnType={userClassRef}
+        >
+          <py.StatementList>
+            <py.VariableDeclaration
+              name="user"
+              initializer={<py.ClassInstantiation target={userClassRef} />}
+            />
+            <>return user</>
+          </py.StatementList>
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+    ]);
+
+    // 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([
+      <py.SourceFile path="models.py">
+        <py.ClassDeclaration name="User" refkey={userClassRef} />
+        <py.FunctionDeclaration name="helper" refkey={helperFuncRef}>
+          pass
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+      <py.SourceFile path="service.py">
+        <py.FunctionDeclaration
+          name="process"
+          parameters={[{ name: "user", type: userClassRef }]}
+          returnType="None"
+        >
+          <py.FunctionCallExpression target={helperFuncRef} />
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+    ]);
+
+    // 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([
+      <py.SourceFile path="types.py">
+        <py.ClassDeclaration name="Result" refkey={resultTypeRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="main.py">
+        <py.FunctionDeclaration name="get_result" returnType={resultTypeRef}>
+          pass
+        </py.FunctionDeclaration>
+      </py.SourceFile>,
+    ]);
+
+    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([
+      <py.SourceFile path="types.py">
+        <py.ClassDeclaration name="Config" refkey={configTypeRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="main.py">
+        <py.VariableDeclaration name="config" type={configTypeRef} omitNone />
+      </py.SourceFile>,
+    ]);
+
+    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([
+      <py.SourceFile path="types.py">
+        <py.ClassDeclaration name="MyType" refkey={myTypeRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="main.py">
+        <py.VariableDeclaration
+          name="value"
+          type={<py.TypeReference refkey={myTypeRef} />}
+          omitNone
+        />
+      </py.SourceFile>,
+    ]);
+
+    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([
+      <py.SourceFile path="base.py">
+        <py.ClassDeclaration name="BaseClass" refkey={baseClassRef} />
+      </py.SourceFile>,
+      <py.SourceFile path="derived.py">
+        <py.ClassDeclaration name="DerivedClass" bases={[baseClassRef]} />
+      </py.SourceFile>,
+    ]);
+
+    // 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(
+      [
+        <py.SourceFile path="models.py">
+          <py.ClassDeclaration name="User" refkey={userClassRef} />
+        </py.SourceFile>,
+        <py.SourceFile path="service.py">
+          <py.FunctionDeclaration
+            name="get_users"
+            parameters={[{ name: "user", type: userClassRef }]}
+            returnType="str"
+          >
+            <py.StatementList>
+              {/* Use cast as a value (function call) to make it a regular import */}
+              <>response = {requestsModule["."].get}("https://example.com")</>
+              <>return {typingModule["."].cast}(str, user)</>
+            </py.StatementList>
+          </py.FunctionDeclaration>
+        </py.SourceFile>,
+      ],
+      { 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)
+
+        `,
+    });
+  });
+});

package/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
       `,

package/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
       `,

package/vitest.config.ts

@@ -2,6 +2,14 @@ import alloyPlugin from "@alloy-js/rollup-plugin";
 import { defineConfig } from "vitest/config";
 
 export default defineConfig({
+  resolve: {
+    conditions: ["source"],
+  },
+  ssr: {
+    resolve: {
+      conditions: ["source"],
+    },
+  },
   esbuild: {
     jsx: "preserve",
     sourcemap: "both",