Add support for extra imports (#342)

There are certain scenarios in Python where we need to add imports at
the top of the file without them necessarily being triggered by a symbol
being referenced. One example is `from __future__ import annotations`.

To allow that, we are adding the `futureImports` prop, which already
existed and wasn't being used.

Also, we are making the module docstring to be rendered at the top, as
well as properly formatting the header and headerComment.

Finally, we are fixing some spacing-related situations to better adhere
to PEP8.

With those changes, this is how a rendered file would look:

```
<module docstring>
<manually inserted imports>
<imports generated by the import engine>
<children>
```

Author: mauriciogardini

.chronus/changes/feature-extra-import-2025-11-5-23-27-25.md

@@ -0,0 +1,7 @@
+---
+changeKind: feature
+packages:
+  - "@alloy-js/python"
+---
+
+Add support for extra imports

packages/python/src/components/FutureStatement.tsx

@@ -0,0 +1,38 @@
+import { type Children } from "@alloy-js/core";
+
+export interface FutureStatementProps {
+  /**
+   * The name of the feature to import from __future__.
+   */
+  feature: string;
+}
+
+/**
+ * A future statement that imports features from __future__.
+ *
+ * Future statements are directives to the compiler that a particular module
+ * should be compiled using syntax or semantics from a future Python release.
+ * They must appear near the top of the module, after the module docstring (if any).
+ *
+ * Use this in the `futureImports` prop of SourceFile to ensure proper placement.
+ *
+ * @example
+ * ```tsx
+ * <SourceFile path="models.py" futureImports={<FutureStatement feature="annotations" />}>
+ *   <ClassDeclaration name="User">
+ *     <PropertyDeclaration name="manager" type="User" />
+ *   </ClassDeclaration>
+ * </SourceFile>
+ * ```
+ * renders to
+ * ```py
+ * from __future__ import annotations
+ *
+ *
+ * class User:
+ *     manager: User
+ * ```
+ */
+export function FutureStatement(props: FutureStatementProps): Children {
+  return <>from __future__ import {props.feature}</>;
+}

packages/python/src/components/SourceFile.tsx

@@ -1,7 +1,9 @@
 import {
+  childrenArray,
   ComponentContext,
   SourceFile as CoreSourceFile,
   createNamedContext,
+  isComponentCreator,
   List,
   Scope,
   Show,
@@ -12,8 +14,51 @@ import {
 import { join } from "pathe";
 import { PythonModuleScope } from "../symbols/index.js";
 import { ImportStatements } from "./ImportStatement.js";
+import { SimpleCommentBlock } from "./PyDoc.js";
 import { Reference } from "./Reference.js";
 
+// Non top-level definitions
+const NON_DEFINITION_NAMES = new Set([
+  "VariableDeclaration",
+  "MemberExpression",
+  "FunctionCallExpression",
+  "ClassInstantiation",
+  "Reference",
+]);
+
+// Wrapper components that we should look inside to find the actual first child
+const WRAPPER_COMPONENT_NAMES = new Set(["StatementList"]);
+
+/**
+ * Checks if the first child is a top-level definition (function or class).
+ * PEP 8 requires 2 blank lines before top-level function and class definitions,
+ * but not before other statements like variable assignments.
+ *
+ * Returns true only if there are children and the first child is a definition.
+ */
+function firstChildIsDefinition(children: Children | undefined): boolean {
+  if (!children) return false;
+  const arr = childrenArray(() => children);
+  if (arr.length === 0) return false;
+  const first = arr[0];
+
+  // Non-component children (strings, numbers, refkeys, etc.) are not definitions
+  if (!isComponentCreator(first)) {
+    return false;
+  }
+
+  const name = first.component.name;
+  if (NON_DEFINITION_NAMES.has(name)) {
+    return false;
+  }
+  // Look inside wrapper components
+  if (WRAPPER_COMPONENT_NAMES.has(name) && first.props?.children) {
+    return firstChildIsDefinition(first.props.children as Children);
+  }
+  // If we get here, it's likely a definition (FunctionDeclaration, ClassDeclaration, etc.)
+  return true;
+}
+
 export interface PythonSourceFileContext {
   scope: PythonModuleScope;
   /** The module name for this file, e.g. 'test' for test.py */
@@ -37,17 +82,23 @@ export interface SourceFileProps {
    */
   children?: Children;
   /**
-   * Header comment to add to the file, which will be rendered at the top of the file.
+   * Content to render at the very top of the file, before everything else.
+   * Use this for shebang lines, encoding declarations, or license headers.
    */
   header?: Children;
   /**
-   * Comment to add to the header, which will be rendered as a comment in the file.
+   * Comment to add at the top of the file, rendered as a Python comment block.
+   * This is a convenience prop for adding copyright notices or other comments.
    */
   headerComment?: string;
   /**
    * Documentation for this module, which will be rendered as a module-level docstring.
    */
   doc?: Children;
+  /**
+   * __future__ imports to render after the docstring but before regular imports.
+   */
+  futureImports?: Children;
 }
 
 /**
@@ -100,21 +151,90 @@ export function SourceFile(props: SourceFileProps) {
     module: path,
   };
 
+  // Check if there are any children
+  const hasChildren =
+    props.children !== undefined &&
+    childrenArray(() => props.children).length > 0;
+
+  // PEP 8 requires 2 blank lines before top-level function/class definitions
+  const needsExtraSpacing = firstChildIsDefinition(props.children);
+
+  // Check if there's any preamble content (header, doc, imports, etc.)
+  const hasPreamble =
+    props.header !== undefined ||
+    props.headerComment !== undefined ||
+    props.doc !== undefined ||
+    props.futureImports !== undefined;
+
   return (
-    <CoreSourceFile path={props.path} filetype="py" reference={Reference}>
-      <Show when={scope.importedModules.size > 0}>
-        <ImportStatements records={scope.importedModules} />
-        <hbr />
+    <CoreSourceFile
+      path={props.path}
+      filetype="py"
+      reference={Reference}
+      header={props.header}
+    >
+      {/* Extra blank line after header when followed by doc/futureImports/children (not headerComment) */}
+      <Show
+        when={
+          props.header !== undefined &&
+          props.headerComment === undefined &&
+          (props.doc !== undefined ||
+            props.futureImports !== undefined ||
+            hasChildren)
+        }
+      >
         <hbr />
       </Show>
+      <Show when={props.headerComment !== undefined}>
+        <SimpleCommentBlock>{props.headerComment}</SimpleCommentBlock>
+        {/* When followed by doc: just newline (no blank line) */}
+        <Show when={props.doc !== undefined}>
+          <hbr />
+        </Show>
+        {/* When followed by futureImports or children directly (no doc): blank line */}
+        <Show
+          when={
+            props.doc === undefined &&
+            (props.futureImports !== undefined || hasChildren)
+          }
+        >
+          <hbr />
+          <hbr />
+        </Show>
+      </Show>
       <Show when={props.doc !== undefined}>
         {props.doc}
-        <hbr />
+        <Show when={props.futureImports !== undefined || hasChildren}>
+          <hbr />
+        </Show>
+      </Show>
+      <Show when={props.futureImports !== undefined}>
+        {props.futureImports}
+        <Show when={hasChildren}>
+          <hbr />
+          <hbr />
+        </Show>
+      </Show>
+      <Show when={scope.importedModules.size > 0}>
+        <ImportStatements records={scope.importedModules} />
+        <Show when={hasChildren}>
+          <hbr />
+          <hbr />
+        </Show>
+      </Show>
+      {/* Extra blank line before top-level definitions */}
+      <Show
+        when={
+          needsExtraSpacing && (hasPreamble || scope.importedModules.size > 0)
+        }
+      >
         <hbr />
       </Show>
       <PythonSourceFileContext.Provider value={sfContext}>
         <Scope value={scope}>
-          <List doubleHardline>{props.children}</List>
+          <Show when={hasChildren}>
+            <List doubleHardline>{props.children}</List>
+          </Show>
         </Scope>
       </PythonSourceFileContext.Provider>
     </CoreSourceFile>

packages/python/src/components/index.ts

@@ -12,6 +12,7 @@ export * from "./EnumMember.js";
 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";

packages/python/test/classdeclarations.test.tsx

@@ -108,6 +108,7 @@ describe("Python Class", () => {
     const mod2Expected = d`
       from mod1 import A
 
+
       class B(A):
           pass
 
@@ -116,6 +117,7 @@ describe("Python Class", () => {
     const mod3Expected = d`
       from folder.mod2 import B
 
+
       class C(B):
           pass
 

packages/python/test/dataclassdeclarations.test.tsx

@@ -27,6 +27,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass
         class User:
             """
@@ -74,6 +75,7 @@ describe("DataclassDeclaration", () => {
         from dataclasses import dataclass
         from dataclasses import KW_ONLY
 
+
         @dataclass
         class User:
             id: int
@@ -106,6 +108,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(frozen=True, slots=True, kw_only=True)
         class User:
             id: int
@@ -141,6 +144,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(init=True, repr=False, eq=True, order=False, unsafe_hash=True, frozen=True, match_args=False, kw_only=True, slots=True, weakref_slot=False)
         class User:
             pass
@@ -178,6 +182,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(slots=True, weakref_slot=True)
         class User:
             pass
@@ -213,6 +218,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(order=True)
         class User:
             pass
@@ -398,6 +404,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(kw_only=True)
         class User:
             id: int
@@ -421,6 +428,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass
         class User(Base):
             pass
@@ -503,6 +511,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(unsafe_hash=True)
         class User:
             pass
@@ -556,6 +565,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(frozen=True)
         class User:
             pass
@@ -578,6 +588,7 @@ describe("DataclassDeclaration", () => {
       d`
         from dataclasses import dataclass
 
+
         @dataclass(slots=True)
         class User:
             pass
@@ -627,6 +638,7 @@ describe("DataclassDeclaration", () => {
       "models.py": `
         from dataclasses import dataclass
 
+
         @dataclass
         class User:
             id: int
@@ -636,6 +648,7 @@ describe("DataclassDeclaration", () => {
       "services.py": `
         from models import User
 
+
         def get_user() -> User:
             user: User = User(1, "Alice")
             return user