fix(python): Handle Duplicative Import Names (#5377)

* Get it working

* Optimize and simplify code

* Address nits

* Add util for generating python class names

* Integrate and update tests

* Update test

* Use named interface for UniqueReferenceValue

* Rename to getFullyQualifiedModulePath

* Add test case

generators/python-v2/ast/src/PythonFile.ts

@@ -2,8 +2,19 @@ import { AstNode } from "./core/AstNode";
 import { Comment } from "./Comment";
 import { Writer } from "./core/Writer";
 import { Reference } from "./Reference";
-import { ModulePath } from "./core/types";
+import { ImportedName, ModulePath } from "./core/types";
 import { StarImport } from "./StarImport";
+import { Class } from "./Class";
+import { Method } from "./Method";
+import { Field } from "./Field";
+import { Type } from "./Type";
+import { createPythonClassName } from "./core/utils";
+
+interface UniqueReferenceValue {
+    modulePath: ModulePath;
+    references: Reference[];
+    referenceNames: Set<string>;
+}
 
 export declare namespace PythonFile {
     interface Args {
@@ -44,47 +55,95 @@ export class PythonFile extends AstNode {
     }
 
     public write(writer: Writer): void {
+        const uniqueReferences = this.deduplicateReferences();
+
+        this.updateWriterRefNameOverrides({ writer, uniqueReferences });
+
         this.writeComments(writer);
-        this.writeImports(writer);
+        this.writeImports({ writer, uniqueReferences });
         this.statements.forEach((statement, idx) => {
             statement.write(writer);
             writer.newLine();
             if (idx < this.statements.length - 1) {
                 writer.newLine();
             }
         });
+
+        writer.unsetRefNameOverrides();
     }
 
     /*******************************
      * Helper Methods
      *******************************/
 
-    private writeComments(writer: Writer): void {
-        this.comments.forEach((comment) => {
-            comment.write(writer);
+    private updateWriterRefNameOverrides({
+        writer,
+        uniqueReferences
+    }: {
+        writer: Writer;
+        uniqueReferences: Map<string, UniqueReferenceValue>;
+    }): void {
+        const references: Reference[] = Array.from(uniqueReferences.values()).flatMap(({ references }) => references);
+
+        // Build up a map of refs to their name overrides, keeping track of names that have been used as we go.
+        const completeRefPathsToNameOverrides: Record<string, ImportedName> = {};
+        const usedNames = this.getInitialUsedNames();
+
+        references.forEach((reference) => {
+            const name = reference.alias ?? reference.name;
+            const fullyQualifiedModulePath = reference.getFullyQualifiedModulePath();
+
+            let nameOverride = name;
+            let modulePathIdx = reference.modulePath.length - 1;
+            let isAlias = !!reference.alias;
+
+            while (usedNames.has(nameOverride)) {
+                isAlias = true;
+
+                const module = reference.modulePath[modulePathIdx];
+                if (modulePathIdx < 0 || !module) {
+                    nameOverride = `_${nameOverride}`;
+                } else {
+                    nameOverride = `${createPythonClassName(module)}${nameOverride}`;
+                }
+
+                modulePathIdx--;
+            }
+            usedNames.add(nameOverride);
+
+            completeRefPathsToNameOverrides[fullyQualifiedModulePath] = {
+                name: nameOverride,
+                isAlias
+            };
         });
 
-        if (this.comments.length > 0) {
-            writer.newLine();
-        }
+        writer.setRefNameOverrides(completeRefPathsToNameOverrides);
     }
 
-    private getImportName(reference: Reference): string {
-        const name = reference.name;
-        const alias = reference.alias;
-        return `${name}${alias ? ` as ${alias}` : ""}`;
+    private getInitialUsedNames(): Set<string> {
+        const usedNames = new Set<string>();
+
+        this.statements.forEach((statement) => {
+            if (statement instanceof Class) {
+                usedNames.add(statement.name);
+            } else if (statement instanceof Method) {
+                usedNames.add(statement.name);
+            } else if (statement instanceof Field) {
+                usedNames.add(statement.name);
+            }
+        });
+
+        return usedNames;
     }
 
-    private writeImports(writer: Writer): void {
+    private deduplicateReferences() {
         // Deduplicate references by their fully qualified paths
-        const uniqueReferences = new Map<
-            string,
-            { modulePath: ModulePath; references: Reference[]; referenceNames: Set<string> }
-        >();
+        const uniqueReferences = new Map<string, UniqueReferenceValue>();
         for (const reference of this.references) {
-            const fullyQualifiedPath = reference.getFullyQualifiedModulePath();
-            const existingRefs = uniqueReferences.get(fullyQualifiedPath);
             const referenceName = reference.name;
+            const fullyQualifiedPath = reference.getFullyQualifiedPath();
+            const existingRefs = uniqueReferences.get(fullyQualifiedPath);
+
             if (existingRefs) {
                 if (!existingRefs.referenceNames.has(referenceName)) {
                     existingRefs.references.push(reference);
@@ -99,6 +158,35 @@ export class PythonFile extends AstNode {
             }
         }
 
+        return uniqueReferences;
+    }
+
+    private writeComments(writer: Writer): void {
+        this.comments.forEach((comment) => {
+            comment.write(writer);
+        });
+
+        if (this.comments.length > 0) {
+            writer.newLine();
+        }
+    }
+
+    private getImportName({ writer, reference }: { writer: Writer; reference: Reference }): string {
+        const nameOverride = writer.getRefNameOverride(reference);
+
+        const name = reference.name;
+        const alias = nameOverride.isAlias ? nameOverride.name : undefined;
+
+        return `${name}${alias ? ` as ${alias}` : ""}`;
+    }
+
+    private writeImports({
+        writer,
+        uniqueReferences
+    }: {
+        writer: Writer;
+        uniqueReferences: Map<string, { modulePath: ModulePath; references: Reference[] }>;
+    }): void {
         for (const [fullyQualifiedPath, { modulePath, references }] of uniqueReferences) {
             const refModulePath = modulePath;
             if (refModulePath[0] === this.path[0]) {
@@ -127,12 +215,16 @@ export class PythonFile extends AstNode {
 
                 // Write the relative import statement
                 writer.write(
-                    `from ${relativePath} import ${references.map((ref) => this.getImportName(ref)).join(", ")}`
+                    `from ${relativePath} import ${references
+                        .map((reference) => this.getImportName({ writer, reference }))
+                        .join(", ")}`
                 );
             } else {
                 // Use fully qualified path
                 writer.write(
-                    `from ${fullyQualifiedPath} import ${references.map((ref) => this.getImportName(ref)).join(", ")}`
+                    `from ${fullyQualifiedPath} import ${references
+                        .map((reference) => this.getImportName({ writer, reference }))
+                        .join(", ")}`
                 );
             }
 

generators/python-v2/ast/src/Reference.ts

@@ -43,7 +43,8 @@ export class Reference extends AstNode {
     }
 
     public write(writer: Writer): void {
-        writer.write(this.alias ?? this.name);
+        const nameOverride = writer.getRefNameOverride(this);
+        writer.write(nameOverride.name);
 
         if (this.genericTypes.length > 0) {
             writer.write("[");
@@ -67,7 +68,11 @@ export class Reference extends AstNode {
         }
     }
 
-    public getFullyQualifiedModulePath(): string {
+    public getFullyQualifiedPath(): string {
         return this.modulePath.join(".");
     }
+
+    public getFullyQualifiedModulePath(): string {
+        return `${this.getFullyQualifiedPath()}.${this.name}`;
+    }
 }

generators/python-v2/ast/src/__test__/PythonFile.test.ts

@@ -316,4 +316,70 @@ describe("PythonFile", () => {
         expect(await writer.toStringFormatted()).toMatchSnapshot();
         expect(file.getReferences()).toHaveLength(2);
     });
+
+    it("Write duplicative import names", async () => {
+        const file = python.file({ path: ["root"] });
+
+        const local_class = python.class_({
+            name: "Car"
+        });
+
+        local_class.add(
+            python.field({
+                name: "car",
+                initializer: python.instantiateClass({
+                    classReference: python.reference({
+                        modulePath: ["root", "cars"],
+                        name: "Car"
+                    }),
+                    arguments_: []
+                })
+            })
+        );
+
+        local_class.add(
+            python.field({
+                name: "car",
+                initializer: python.instantiateClass({
+                    classReference: python.reference({
+                        modulePath: ["root", "transportation", "vehicles"],
+                        name: "Car"
+                    }),
+                    arguments_: []
+                })
+            })
+        );
+
+        local_class.add(
+            python.field({
+                name: "automobile",
+                initializer: python.instantiateClass({
+                    classReference: python.reference({
+                        modulePath: ["root", "automobiles"],
+                        name: "Car"
+                    }),
+                    arguments_: []
+                })
+            })
+        );
+
+        local_class.add(
+            python.field({
+                name: "vehicle",
+                initializer: python.instantiateClass({
+                    classReference: python.reference({
+                        modulePath: ["root", "vehicles", "automobiles"],
+                        name: "Car"
+                    }),
+                    arguments_: []
+                })
+            })
+        );
+
+        file.addStatement(local_class);
+
+        file.write(writer);
+        expect(await writer.toStringFormatted()).toMatchSnapshot();
+        expect(file.getReferences()).toHaveLength(4);
+    });
 });

generators/python-v2/ast/src/__test__/__snapshots__/PythonFile.test.ts.snap

@@ -112,6 +112,21 @@ my_variable: ImportedClass.nested.attribute
 "
 `;
 
+exports[`PythonFile > Write duplicative import names 1`] = `
+"from .cars import Car as CarsCar
+from .transportation.vehicles import Car as VehiclesCar
+from .automobiles import Car as AutomobilesCar
+from .vehicles.automobiles import Car as VehiclesAutomobilesCar
+
+
+class Car:
+    car = CarsCar()
+    car = VehiclesCar()
+    automobile = AutomobilesCar()
+    vehicle = VehiclesAutomobilesCar()
+"
+`;
+
 exports[`PythonFile > Write star imports 1`] = `
 "# flake8: noqa: F401, F403
 

generators/python-v2/ast/src/__test__/core/utils.test.ts

@@ -0,0 +1,33 @@
+import { createPythonClassName } from "../../core/utils";
+
+describe("Casing", () => {
+    describe("createPythonClassName", () => {
+        const testCases: [string, string][] = [
+            // Basic cases
+            ["hello world", "HelloWorld"],
+            ["simpleTestCase", "SimpleTestCase"],
+            // Special characters
+            ["hello-world", "HelloWorld"],
+            ["$special#characters%", "SpecialCharacters"],
+            // Numbers
+            ["123 invalid class name", "Class123InvalidClassName"],
+            ["mixed 123 cases", "Mixed123Cases"],
+            // Underscores
+            ["_leading_underscores_", "LeadingUnderscores"],
+            ["trailing_underscores_", "TrailingUnderscores"],
+            ["_123numbers_starting", "Class123NumbersStarting"],
+            // Empty and invalid input
+            ["", "Class"],
+            ["123", "Class123"],
+            ["_123_", "Class123"],
+            // Complex cases
+            ["complex mix_of-DifferentCases", "ComplexMixOfDifferentCases"],
+            ["ALLCAPS input", "ALLCAPSInput"], // Preserve ALLCAPS as requested
+            ["PascalCaseAlready", "PascalCaseAlready"]
+        ];
+
+        it.each<[string, string]>(testCases)("should convert %s' to %s'", (input, expected) => {
+            expect(createPythonClassName(input)).toBe(expected);
+        });
+    });
+});

generators/python-v2/ast/src/core/Writer.ts

@@ -1,9 +1,32 @@
 import { AbstractWriter } from "@fern-api/base-generator";
 import { Config } from "@wasm-fmt/ruff_fmt";
+import { Reference } from "../Reference";
+import { ImportedName } from "./types";
 
 export declare namespace Writer {}
 
 export class Writer extends AbstractWriter {
+    private fullyQualifiedModulePathsToImportedNames: Record<string, ImportedName> = {};
+
+    public setRefNameOverrides(completeRefPathsToNameOverrides: Record<string, ImportedName>): void {
+        this.fullyQualifiedModulePathsToImportedNames = completeRefPathsToNameOverrides;
+    }
+
+    public unsetRefNameOverrides(): void {
+        this.fullyQualifiedModulePathsToImportedNames = {};
+    }
+
+    public getRefNameOverride(reference: Reference): ImportedName {
+        const explicitNameOverride =
+            this.fullyQualifiedModulePathsToImportedNames[reference.getFullyQualifiedModulePath()];
+
+        if (explicitNameOverride) {
+            return explicitNameOverride;
+        }
+
+        return { name: reference.alias ?? reference.name, isAlias: !!reference.alias };
+    }
+
     public toString(): string {
         return this.buffer;
     }

generators/python-v2/ast/src/core/types.ts

@@ -1,2 +1,8 @@
 export type ModulePath = string[] | Readonly<string[]>;
+
 export type AttrPath = string[] | Readonly<string[]>;
+
+export interface ImportedName {
+    name: string;
+    isAlias?: boolean;
+}