Add support for reference aliases

fern-api · Oct 17, 2024 · 1d48581 · 1d48581
1 parent 50137a7
commit 1d48581
Show file tree

Hide file tree

Showing 5 changed files with 89 additions and 1 deletion.
diff --git a/generators/pythonv2/codegen/src/ast/PythonFile.ts b/generators/pythonv2/codegen/src/ast/PythonFile.ts
@@ -92,6 +92,10 @@ export class PythonFile extends AstNode {
             } else {
                 // Use fully qualified path
                 writer.write(`from ${reference.getFullyQualifiedModulePath()} import ${reference.getName()}`);
+
+                if (reference.getAlias()) {
+                    writer.write(` as ${reference.getAlias()}`);
+                }
             }
             writer.newLine();
         });

diff --git a/generators/pythonv2/codegen/src/ast/Reference.ts b/generators/pythonv2/codegen/src/ast/Reference.ts
@@ -14,19 +14,23 @@ export declare namespace Reference {
         modulePath?: string[];
         /* The generic types of the reference */
         genericTypes?: Type[];
+        /* The alias of the reference */
+        alias?: string;
     }
 }
 
 export class Reference extends AstNode {
     private name: string;
     private modulePath: string[];
     private genericTypes: Type[];
+    private alias: string | undefined;
 
-    constructor({ name, modulePath, genericTypes }: Reference.Args) {
+    constructor({ name, modulePath, genericTypes, alias }: Reference.Args) {
         super();
         this.name = name;
         this.modulePath = modulePath ?? [];
         this.genericTypes = genericTypes ?? [];
+        this.alias = alias;
     }
 
     public write(writer: Writer): void {
@@ -42,6 +46,11 @@ export class Reference extends AstNode {
             });
             writer.write("]");
         }
+
+        if (this.alias) {
+            writer.write(" as ");
+            writer.write(this.alias);
+        }
     }
 
     public getName(): string {
@@ -59,4 +68,8 @@ export class Reference extends AstNode {
     public getFullyQualifiedName(): string {
         return [this.getFullyQualifiedModulePath(), this.name].join(".");
     }
+
+    public getAlias(): string | undefined {
+        return this.alias;
+    }
 }
diff --git a/generators/pythonv2/codegen/src/ast/__test__/PythonFile.test.ts b/generators/pythonv2/codegen/src/ast/__test__/PythonFile.test.ts
@@ -111,4 +111,50 @@ describe("PythonFile", () => {
         file.write(writer);
         expect(writer.toString()).toMatchSnapshot();
     });
+
+    it("Add a class with an absolute import and alias", () => {
+        const file = python.file({
+            moduleName: "test_module",
+            path: ["test"],
+            name: "test_file"
+        });
+
+        const absoluteRef = python.reference({
+            modulePath: ["external_module", "submodule"],
+            name: "ExternalClass",
+            alias: "AliasedClass"
+        });
+        const testClass = python.class_({
+            name: "TestClassWithAlias",
+            extends_: [absoluteRef]
+        });
+        writer.addReference(absoluteRef);
+        file.addStatement(testClass);
+
+        file.write(writer);
+        expect(writer.toString()).toMatchSnapshot();
+    });
+
+    it("Add a class with a relative import and alias", () => {
+        const file = python.file({
+            moduleName: "test_module",
+            path: ["test", "subdir"],
+            name: "test_file"
+        });
+
+        const relativeRef = python.reference({
+            modulePath: ["test_module", "test", "sibling_dir"],
+            name: "SiblingClass",
+            alias: "AliasedSibling"
+        });
+        const testClass = python.class_({
+            name: "TestClassWithRelativeAlias",
+            extends_: [relativeRef]
+        });
+        writer.addReference(relativeRef);
+        file.addStatement(testClass);
+
+        file.write(writer);
+        expect(writer.toString()).toMatchSnapshot();
+    });
 });
diff --git a/generators/pythonv2/codegen/src/ast/__test__/Reference.test.ts b/generators/pythonv2/codegen/src/ast/__test__/Reference.test.ts
@@ -49,5 +49,14 @@ describe("Reference", () => {
             });
             expect(reference.toString()).toBe("DoubleGenericClass[str, int]");
         });
+
+        it("handles class with alias", () => {
+            const reference = python.reference({
+                name: "AliasClass",
+                modulePath: ["module"],
+                alias: "Alias"
+            });
+            expect(reference.toString()).toBe("AliasClass as Alias");
+        });
     });
 });
diff --git a/generators/pythonv2/codegen/src/ast/__test__/__snapshots__/PythonFile.test.ts.snap b/generators/pythonv2/codegen/src/ast/__test__/__snapshots__/PythonFile.test.ts.snap
@@ -27,6 +27,22 @@ class DeeplyNestedTestClass(DeepClass):
 "
 `;
 
+exports[`PythonFile > Add a class with a relative import and alias 1`] = `
+"from ...test_module.test.sibling_dir import SiblingClass
+
+class TestClassWithRelativeAlias(SiblingClass as AliasedSibling):
+    pass
+"
+`;
+
+exports[`PythonFile > Add a class with an absolute import and alias 1`] = `
+"from external_module.submodule import ExternalClass as AliasedClass
+
+class TestClassWithAlias(ExternalClass as AliasedClass):
+    pass
+"
+`;
+
 exports[`PythonFile > Add a class with no references 1`] = `
 "class TestClass:
     pass