feat(python): Introduce Type, Annotation, Variable, and Integrate w/ Class

generators/pythonv2/codegen/package.json

@@ -27,6 +27,7 @@
     "depcheck": "depcheck"
   },
   "dependencies": {
+    "@fern-api/core-utils": "workspace:*",
     "@fern-api/generator-commons": "workspace:*",
     "@fern-fern/ir-sdk": "^53.7.0"
   },

generators/pythonv2/codegen/src/ast/Annotation.ts

@@ -0,0 +1,27 @@
+import { AstNode } from "./core/AstNode";
+import { Writer } from "./core/Writer";
+
+export declare namespace Annotation {
+    interface Args {
+        /* The type hint */
+        type: string | AstNode;
+    }
+}
+
+export class Annotation extends AstNode {
+    private type: string | AstNode;
+
+    constructor(args: Annotation.Args) {
+        super();
+        this.type = args.type;
+    }
+
+    public write(writer: Writer): void {
+        writer.write(": ");
+        if (typeof this.type === "string") {
+            writer.write(this.type);
+        } else {
+            this.type.write(writer);
+        }
+    }
+}

generators/pythonv2/codegen/src/ast/Class.ts

@@ -1,5 +1,6 @@
 import { AstNode } from "./core/AstNode";
 import { Writer } from "./core/Writer";
+import { Variable } from "./Variable";
 
 export declare namespace Class {
     interface Args {
@@ -11,12 +12,34 @@ export declare namespace Class {
 export class Class extends AstNode {
     public readonly name: string;
 
+    private variables: Variable[] = [];
+
     constructor({ name }: Class.Args) {
         super();
         this.name = name;
     }
 
     public write(writer: Writer): void {
         writer.write(`class ${this.name}:`);
+        writer.newLine();
+
+        writer.indent();
+        this.writeVariables({ writer });
+        writer.dedent();
+    }
+
+    public addVariable(variable: Variable): void {
+        this.variables.push(variable);
+    }
+
+    private writeVariables({ writer }: { writer: Writer }): void {
+        this.variables.forEach((variable, index) => {
+            variable.write(writer);
+            writer.writeNewLineIfLastLineNot();
+
+            if (index < this.variables.length - 1) {
+                writer.newLine();
+            }
+        });
     }
 }

generators/pythonv2/codegen/src/ast/Type.ts

@@ -0,0 +1,195 @@
+import { assertNever } from "@fern-api/core-utils";
+import { AstNode } from "./core/AstNode";
+import { Writer } from "./core/Writer";
+
+type InternalType = Int | Float | Bool | Str | Bytes | List | Set | Tuple | Dict | None | Optional | Union | Any;
+
+interface Int {
+    type: "int";
+}
+
+interface Float {
+    type: "float";
+}
+
+interface Bool {
+    type: "bool";
+}
+
+interface Str {
+    type: "str";
+}
+
+interface Bytes {
+    type: "bytes";
+}
+
+interface List {
+    type: "list";
+    value: Type;
+}
+
+interface Set {
+    type: "set";
+    value: Type;
+}
+
+interface Tuple {
+    type: "tuple";
+    values: Type[];
+}
+
+interface Dict {
+    type: "dict";
+    keyType: Type;
+    valueType: Type;
+}
+
+interface None {
+    type: "none";
+}
+
+interface Optional {
+    type: "optional";
+    value: Type;
+}
+
+interface Union {
+    type: "union";
+    values: Type[];
+}
+
+interface Any {
+    type: "any";
+}
+
+export class Type extends AstNode {
+    private internalType: InternalType;
+
+    private constructor(internalType: InternalType) {
+        super();
+        this.internalType = internalType;
+    }
+
+    public static int(): Type {
+        return new Type({ type: "int" });
+    }
+
+    public static float(): Type {
+        return new Type({ type: "float" });
+    }
+
+    public static bool(): Type {
+        return new Type({ type: "bool" });
+    }
+
+    public static str(): Type {
+        return new Type({ type: "str" });
+    }
+
+    public static bytes(): Type {
+        return new Type({ type: "bytes" });
+    }
+
+    public static list(value: Type): Type {
+        return new Type({ type: "list", value });
+    }
+
+    public static set(value: Type): Type {
+        return new Type({ type: "set", value });
+    }
+
+    public static tuple(values: Type[]): Type {
+        return new Type({ type: "tuple", values });
+    }
+
+    public static dict(keyType: Type, valueType: Type): Type {
+        return new Type({ type: "dict", keyType, valueType });
+    }
+
+    public static none(): Type {
+        return new Type({ type: "none" });
+    }
+
+    public static optional(value: Type): Type {
+        return new Type({ type: "optional", value });
+    }
+
+    public static union(values: Type[]): Type {
+        return new Type({ type: "union", values });
+    }
+
+    public static any(): Type {
+        return new Type({ type: "any" });
+    }
+
+    public write(writer: Writer): void {
+        switch (this.internalType.type) {
+            case "int":
+                writer.write("int");
+                break;
+            case "float":
+                writer.write("float");
+                break;
+            case "bool":
+                writer.write("bool");
+                break;
+            case "str":
+                writer.write("str");
+                break;
+            case "bytes":
+                writer.write("bytes");
+                break;
+            case "list":
+                writer.write("List[");
+                this.internalType.value.write(writer);
+                writer.write("]");
+                break;
+            case "set":
+                writer.write("Set[");
+                this.internalType.value.write(writer);
+                writer.write("]");
+                break;
+            case "tuple":
+                writer.write("Tuple[");
+                this.internalType.values.forEach((value, index) => {
+                    if (index > 0) {
+                        writer.write(", ");
+                    }
+                    value.write(writer);
+                });
+                writer.write("]");
+                break;
+            case "dict":
+                writer.write("Dict[");
+                this.internalType.keyType.write(writer);
+                writer.write(", ");
+                this.internalType.valueType.write(writer);
+                writer.write("]");
+                break;
+            case "none":
+                writer.write("None");
+                break;
+            case "optional":
+                writer.write("Optional[");
+                this.internalType.value.write(writer);
+                writer.write("]");
+                break;
+            case "union":
+                writer.write("Union[");
+                this.internalType.values.forEach((value, index) => {
+                    if (index > 0) {
+                        writer.write(", ");
+                    }
+                    value.write(writer);
+                });
+                writer.write("]");
+                break;
+            case "any":
+                writer.write("Any");
+                break;
+            default:
+                assertNever(this.internalType);
+        }
+    }
+}

generators/pythonv2/codegen/src/ast/Variable.ts

@@ -0,0 +1,37 @@
+import { AstNode } from "./core/AstNode";
+import { Writer } from "./core/Writer";
+import { Annotation } from "./Annotation";
+
+export declare namespace Variable {
+    interface Args {
+        /* The name of the variable */
+        name: string;
+        /* The type annotation of the variable */
+        type?: Annotation;
+        /* The initializer for the variable */
+        initializer?: string;
+    }
+}
+
+export class Variable extends AstNode {
+    public readonly name: string;
+    public readonly type: Annotation | undefined;
+    public readonly initializer: string | undefined;
+
+    constructor({ name, type, initializer }: Variable.Args) {
+        super();
+        this.name = name;
+        this.type = type;
+        this.initializer = initializer;
+    }
+
+    public write(writer: Writer): void {
+        writer.write(this.name);
+        if (this.type) {
+            this.type.write(writer);
+        }
+        if (this.initializer != null) {
+            writer.write(` = ${this.initializer}`);
+        }
+    }
+}

generators/pythonv2/codegen/src/ast/__test__/Class.test.ts

@@ -7,4 +7,14 @@ describe("class", () => {
         });
         expect(clazz.toString()).toMatchSnapshot();
     });
+
+    it("variables with annotation and initializer", async () => {
+        const clazz = python.class_({
+            name: "Car"
+        });
+        clazz.addVariable(
+            python.variable({ name: "color", type: python.annotation({ type: "str" }), initializer: "'red'" })
+        );
+        expect(clazz.toString()).toMatchSnapshot();
+    });
 });