Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(python): Support Adding Statements to Classes and Methods #4949

Merged
merged 3 commits into from
Oct 18, 2024
Merged

feat(python): Support Adding Statements to Classes and Methods #4949

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6b03c42
Support arbitrary statements in class
noanflaherty Oct 18, 2024
1b427b5
Add support for arbitrary statements in methods
noanflaherty Oct 18, 2024
59627f2
Merge branch 'main' into noa/support-arbitrary-statements
noanflaherty Oct 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 12 additions & 21 deletions generators/pythonv2/codegen/src/ast/Class.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,7 @@ export class Class extends AstNode {
public readonly name: string;
public readonly extends_: Reference[];
public readonly decorators: Decorator[];
private fields: Field[] = [];
private body?: CodeBlock;
private statements: AstNode[] = [];

constructor({ name, extends_, decorators }: Class.Args) {
super();
Expand Down Expand Up @@ -55,32 +54,24 @@ export class Class extends AstNode {
writer.newLine();

writer.indent();
let hasContents = false;
if (this.fields.length) {
this.writeFields({ writer });
hasContents = true;
}
if (this.body) {
this.body.write(writer);
hasContents = true;
}
if (!hasContents) {
if (this.statements.length) {
this.writeStatements({ writer });
} else {
writer.write("pass");
}
writer.dedent();
}

public addField(field: Field): void {
this.fields.push(field);
}
public addStatement(statement: AstNode): void {
this.statements.push(statement);

public addBody(body: CodeBlock): void {
this.body = body;
statement.getReferences().forEach((reference) => {
this.addReference(reference);
});
}

private writeFields({ writer }: { writer: Writer }): void {
this.fields.forEach((field, index) => {
field.write(writer);
private writeStatements({ writer }: { writer: Writer }): void {
this.statements.forEach((statement, index) => {
statement.write(writer);
writer.writeNewLineIfLastLineNot();
});
}
Expand Down
26 changes: 18 additions & 8 deletions generators/pythonv2/codegen/src/ast/Method.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ export declare namespace Method {
parameters: Parameter[];
/* The return type of the method */
return_?: Type;
/* The body of the method */
body?: CodeBlock;
/* The docstring for the method */
docstring?: string;
/* The type of the method if defined within the context of a class */
Expand All @@ -37,18 +35,17 @@ export declare namespace Method {
export class Method extends AstNode {
public readonly name: string;
public readonly return: Type | undefined;
public readonly body: CodeBlock | undefined;
public readonly docstring: string | undefined;
public readonly type: ClassMethodType | undefined;
private readonly parameters: Parameter[];
private readonly decorators: Decorator[];
private readonly statements: AstNode[] = [];

constructor({ name, parameters, return_, body, docstring, type, decorators }: Method.Args) {
constructor({ name, parameters, return_, docstring, type, decorators }: Method.Args) {
super();
this.name = name;
this.parameters = parameters;
this.return = return_;
this.body = body;
this.docstring = docstring;
this.type = type;
this.decorators = decorators ?? [];
Expand All @@ -58,6 +55,14 @@ export class Method extends AstNode {
return this.name;
}

public addStatement(statement: AstNode): void {
this.statements.push(statement);

statement.getReferences().forEach((reference) => {
this.addReference(reference);
});
}

public write(writer: Writer): void {
// Write decorators
this.decorators.forEach((decorator) => {
Expand Down Expand Up @@ -119,15 +124,20 @@ export class Method extends AstNode {
}

// Write method body
if (this.body) {
if (this.statements.length) {
writer.indent();
this.body.write(writer);
this.statements.forEach((statement, index) => {
statement.write(writer);
if (index < this.statements.length - 1) {
writer.newLine();
}
});
writer.dedent();
} else {
writer.indent();
writer.write("pass");
writer.newLine();
writer.dedent();
}
writer.newLine();
}
}
31 changes: 29 additions & 2 deletions generators/pythonv2/codegen/src/ast/__test__/Class.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,10 @@ describe("class", () => {
const clazz = python.class_({
name: "Car"
});
clazz.addField(
clazz.addStatement(
python.field({ name: "color", type: python.Type.str(), initializer: python.codeBlock("'red'") })
);
clazz.addField(
clazz.addStatement(
python.field({
name: "partNameById",
type: python.Type.dict(python.Type.int(), python.Type.str()),
Expand Down Expand Up @@ -78,4 +78,31 @@ describe("class", () => {
clazz.write(writer);
expect(await writer.toStringFormatted()).toMatchSnapshot();
});

it("should generate a class with local classes", async () => {
const clazz = python.class_({
name: "OuterClass"
});

const parentClassRef = python.reference({
name: "ParentClass",
modulePath: ["some_module"]
});

const innerClassDef = python.class_({
name: "InnerClass",
extends_: [parentClassRef]
});
const innerMethod = python.method({
name: "inner_method",
parameters: [python.parameter({ name: "self", type: python.Type.str() })]
});
innerMethod.addStatement(python.codeBlock('return "Inner method called"'));
innerClassDef.addStatement(innerMethod);

clazz.addStatement(innerClassDef);

clazz.write(writer);
expect(await writer.toStringFormatted()).toMatchSnapshot();
});
});
34 changes: 32 additions & 2 deletions generators/pythonv2/codegen/src/ast/__test__/Method.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,9 @@ describe("Method", () => {
it("should generate a method with a body", async () => {
const method = python.method({
name: "with_body",
parameters: [],
body: python.codeBlock("return True")
parameters: []
});
method.addStatement(python.codeBlock("return True"));
method.write(writer);
expect(await writer.toStringFormatted()).toMatchSnapshot();
});
Expand Down Expand Up @@ -133,5 +133,35 @@ describe("Method", () => {
method.write(writer);
expect(await writer.toStringFormatted()).toMatchSnapshot();
});

it("should generate a method with local classes", async () => {
const method = python.method({
name: "method_with_local_classes",
parameters: []
});

const parentClassRef = python.reference({
name: "ParentClass",
modulePath: ["some_module"]
});

const childClassDef = python.class_({
name: "ChildClass",
extends_: [parentClassRef]
});
const childMethod = python.method({
name: "child_method",
parameters: [python.parameter({ name: "self", type: python.Type.str() })]
});
childMethod.addStatement(python.codeBlock("self.parent_method()"));
childMethod.addStatement(python.codeBlock('return "Child method called"'));
childClassDef.addStatement(childMethod);

method.addStatement(childClassDef);
method.addStatement(python.codeBlock("return ChildClass()"));

method.write(writer);
expect(await writer.toStringFormatted()).toMatchSnapshot();
});
});
});
2 changes: 1 addition & 1 deletion generators/pythonv2/codegen/src/ast/__test__/PythonFile.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ describe("PythonFile", () => {
name: "TestClass"
});
testClass.addReference(python.reference({ modulePath: ["itertools"], name: "chain" }));
testClass.addBody(python.codeBlock("flat_list = list(itertools.chain([[1, 2], [3, 4]]))"));
testClass.addStatement(python.codeBlock("flat_list = list(itertools.chain([[1, 2], [3, 4]]))"));

file.addStatement(testClass);

Expand Down
8 changes: 8 additions & 0 deletions generators/pythonv2/codegen/src/ast/__test__/__snapshots__/Class.test.ts.snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,11 @@ exports[`class > inherits from two parent classes 1`] = `
pass
"
`;

exports[`class > should generate a class with local classes 1`] = `
"class OuterClass:
class InnerClass(ParentClass):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[q] where does ParentClass come from? Are we missing an import statement, or is that oos for this PR?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Imports only get resolved within the context of a PythonFile being written.

def inner_method(self: str):
return "Inner method called"
"
`;
11 changes: 11 additions & 0 deletions generators/pythonv2/codegen/src/ast/__test__/__snapshots__/Method.test.ts.snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,17 @@ exports[`Method > toString > should generate a method with a specified return ty
"
`;

exports[`Method > toString > should generate a method with local classes 1`] = `
"def method_with_local_classes():
class ChildClass(ParentClass):
def child_method(self: str):
self.parent_method()
return "Child method called"

return ChildClass()
"
`;

exports[`Method > toString > should generate a method with multiple arguments 1`] = `
"def multi_args(arg1: str, arg2: int):
pass
Expand Down
Loading