fix(python): Fix Reference Tracking for Class and Method AST Nodes (

#5254) * Fix class reference tracking * Fix method reference tracking
fern-api · Nov 22, 2024 · b34f63b · b34f63b
1 parent d594089
commit b34f63b
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 0 deletions.
diff --git a/generators/python-v2/ast/src/Class.ts b/generators/python-v2/ast/src/Class.ts
@@ -35,6 +35,10 @@ export class Class extends AstNode {
         this.extends_.forEach((parentClassReference) => {
             this.inheritReferences(parentClassReference);
         });
+
+        this.decorators.forEach((decorator) => {
+            this.inheritReferences(decorator);
+        });
     }
 
     public write(writer: Writer): void {

diff --git a/generators/python-v2/ast/src/Method.ts b/generators/python-v2/ast/src/Method.ts
@@ -50,6 +50,20 @@ export class Method extends AstNode {
         this.type = type;
         this.decorators = decorators ?? [];
         this.static_ = static_ ?? false;
+
+        this.parameters.forEach((parameter) => {
+            this.inheritReferences(parameter);
+        });
+
+        this.inheritReferences(this.return);
+
+        this.decorators.forEach((decorator) => {
+            this.inheritReferences(decorator);
+        });
+
+        this.statements.forEach((statements) => {
+            this.inheritReferences(statements);
+        });
     }
 
     public addStatement(statement: AstNode): void {

diff --git a/generators/python-v2/ast/src/__test__/Class.test.ts b/generators/python-v2/ast/src/__test__/Class.test.ts
@@ -75,6 +75,7 @@ describe("class", () => {
         });
         clazz.write(writer);
         expect(await writer.toStringFormatted()).toMatchSnapshot();
+        expect(clazz.getReferences().length).toBe(1);
     });
 
     it("should generate a class with local classes", async () => {

diff --git a/generators/python-v2/ast/src/__test__/Method.test.ts b/generators/python-v2/ast/src/__test__/Method.test.ts
@@ -18,6 +18,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate an instance method", async () => {
@@ -28,6 +29,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a class method", async () => {
@@ -38,6 +40,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with one argument", async () => {
@@ -47,6 +50,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with multiple arguments", async () => {
@@ -59,6 +63,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with a specified return type", async () => {
@@ -69,6 +74,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method without a specified return type", async () => {
@@ -78,6 +84,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with a body", async () => {
@@ -88,6 +95,7 @@ describe("Method", () => {
             method.addStatement(python.codeBlock("return True"));
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method without a body", async () => {
@@ -97,6 +105,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with a docstring", async () => {
@@ -107,6 +116,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method without a docstring", async () => {
@@ -116,6 +126,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(0);
         });
 
         it("should generate a method with a decorator", async () => {
@@ -133,6 +144,7 @@ describe("Method", () => {
             });
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(1);
         });
 
         it("should generate a method with local classes", async () => {
@@ -163,6 +175,7 @@ describe("Method", () => {
 
             method.write(writer);
             expect(await writer.toStringFormatted()).toMatchSnapshot();
+            expect(method.getReferences().length).toBe(1);
         });
     });
 });