diff --git a/README.md b/README.md index 53467c1..3cf017a 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,7 @@ This repository contains the following libraries: - Java AST library (java-ast) - C# AST library (csharp-ast) +- Python AST library (python-ast) - Common AST types and interfaces (ast-types) ## Overview @@ -26,7 +27,7 @@ The AST libraries are not intended to cover all language functionality. Instead, When more specialized or custom code is needed, the `CodeBlock` can be used as a generic node that can include any code as a string. This flexibility allows you to generate both structured AST-based code and custom code blocks when necessary. ```typescript -import { CodeBlock } from '@amplication/java-ast'; // or '@amplication/csharp-ast' +import { CodeBlock } from '@amplication/java-ast'; // or '@amplication/csharp-ast' or '@amplication/python-ast' // Create a custom code block when needed const customLogic = new CodeBlock(` @@ -64,6 +65,15 @@ The C# AST library provides functionality for generating C# code through an abst - Attributes and documentation - Namespace management +#### python-ast +The Python AST library provides functionality for generating Python code through an abstract syntax tree. It supports: +- Class and function definitions +- Method decorators and type annotations +- Module and import management +- Docstring generation +- Static and class methods +- Async functions + ## Installation To install the libraries: @@ -75,6 +85,9 @@ npm install @amplication/java-ast # For C# AST npm install @amplication/csharp-ast +# For Python AST +npm install @amplication/python-ast + # For AST Types npm install @amplication/ast-types ``` @@ -87,6 +100,8 @@ To use these libraries in your project: import { Class, Interface, Method } from '@amplication/java-ast'; // or import { Class, Interface, Method } from '@amplication/csharp-ast'; +// or +import { ClassDef, FunctionDef, Decorator } from '@amplication/python-ast'; ``` ## Contributing diff --git a/libs/python-ast/README.md b/libs/python-ast/README.md new file mode 100644 index 0000000..1ec56a8 --- /dev/null +++ b/libs/python-ast/README.md @@ -0,0 +1,307 @@ +# Python AST + +This library provides an Abstract Syntax Tree (AST) representation for Python source code, focusing on the core language features necessary for defining classes, functions, and other declarations, while using a generic `CodeBlock` for unsupported language features. + +## Key Features + +- Core Python language constructs (modules, classes, functions) +- Import management and type annotations +- Generic code block for unsupported language features +- Clean and consistent API that matches other Amplication AST libraries +- Support for static and class methods +- Async function support +- Type hints and annotations + +## Installation + +```bash +npm install @amplication/python-ast +``` + +## Usage + +### Creating a Python Module with Imports + +```typescript +import { + Module, + Import, + ClassReference +} from '@amplication/python-ast'; + +// Create a module +const module = new Module({ + name: 'user_service', +}); + +// Add imports +module.addImport(new Import({ + from: 'typing', + names: ['List', 'Optional'] +})); + +module.addImport(new Import({ + from: 'datetime', + names: ['datetime'] +})); + +// Result: +// from typing import List, Optional +// from datetime import datetime +``` + +### Creating a Complete Python Class + +```typescript +import { + ClassDef, + FunctionDef, + Parameter, + ClassReference, + CodeBlock, + Module, + Decorator, + Return +} from '@amplication/python-ast'; + +// Create a class with inheritance +const userClass = new ClassDef({ + name: 'User', + moduleName: 'models', + docstring: 'Represents a user in the system', + bases: [ + new ClassReference({ name: 'BaseModel', moduleName: 'database.models' }) + ] +}); + +// Add class attributes with type annotations +userClass.addAttribute(new CodeBlock({ + code: 'created_at: datetime = datetime.now()' +})); + +// Add constructor +const initMethod = new FunctionDef({ + name: '__init__', + parameters: [ + new Parameter({ name: 'self' }), + new Parameter({ + name: 'username', + type: new ClassReference({ name: 'str' }) + }), + new Parameter({ + name: 'email', + type: new ClassReference({ name: 'str' }) + }), + new Parameter({ + name: 'age', + type: new ClassReference({ name: 'Optional', genericTypes: [new ClassReference({ name: 'int' })] }) + }) + ], + docstring: 'Initialize a new User instance' +}); + +initMethod.addStatement(new CodeBlock({ + code: 'self.username = username\nself.email = email\nself.age = age' +})); + +userClass.addMethod(initMethod); + +// Add a static method +const createMethod = new FunctionDef({ + name: 'create_user', + isStatic: true, + parameters: [ + new Parameter({ + name: 'username', + type: new ClassReference({ name: 'str' }) + }), + new Parameter({ + name: 'email', + type: new ClassReference({ name: 'str' }) + }) + ], + returnType: new ClassReference({ name: 'User' }), + docstring: 'Create a new user instance' +}); + +createMethod.addStatement(new CodeBlock({ + code: 'user = User(username, email)\nuser.save()\nreturn user' +})); + +userClass.addMethod(createMethod); + +// Add an async method +const fetchDataMethod = new FunctionDef({ + name: 'fetch_data', + isAsync: true, + parameters: [new Parameter({ name: 'self' })], + returnType: new ClassReference({ name: 'dict' }), + docstring: 'Fetch user data asynchronously' +}); + +fetchDataMethod.addStatement(new CodeBlock({ + code: 'data = await api.get_user_data(self.username)\nreturn data' +})); + +userClass.addMethod(fetchDataMethod); + +// Create a module and add the class +const module = new Module({ name: 'models' }); +module.addClass(userClass); + +// This will generate: +/* +from database.models import BaseModel +from datetime import datetime +from typing import Optional + +class User(BaseModel): + """Represents a user in the system""" + + created_at: datetime = datetime.now() + + def __init__(self, username: str, email: str, age: Optional[int]): + """Initialize a new User instance""" + self.username = username + self.email = email + self.age = age + + @staticmethod + def create_user(username: str, email: str) -> "User": + """Create a new user instance""" + user = User(username, email) + user.save() + return user + + async def fetch_data(self) -> dict: + """Fetch user data asynchronously""" + data = await api.get_user_data(self.username) + return data +*/ +``` + +### Using CodeBlock for Unsupported Features + +The `CodeBlock` class is useful for Python features not directly supported by the AST library: + +```typescript +// Exception handling +const tryExceptBlock = new CodeBlock({ + code: ` +try: + result = process_data() + return result +except ValueError as e: + logger.error(f"Invalid data: {e}") + raise +finally: + cleanup_resources() + `, + references: [ + new ClassReference({ name: 'ValueError' }), + new ClassReference({ name: 'logger', moduleName: 'logging' }) + ] +}); + +// Context managers +const withBlock = new CodeBlock({ + code: ` +with open(file_path, 'r') as file: + content = file.read() + process_content(content) + ` +}); + +// Decorators with arguments +const decoratedMethod = new FunctionDef({ + name: 'process_request', + decorators: [ + new Decorator({ + name: 'retry', + arguments: ['max_attempts=3', 'delay=1'], + moduleName: 'utils.decorators' + }) + ] +}); +``` + +## API Reference + +The library provides the following main components: + +- **Module**: Top-level container for Python code + - Manages imports and class definitions + - Handles module-level code organization + +- **ClassDef**: Class definition with methods and attributes + - Supports inheritance + - Manages class attributes and methods + - Handles docstrings and decorators + +- **FunctionDef**: Function or method definition + - Supports static and class methods + - Handles async functions + - Manages parameters and return types + - Supports decorators + +- **Parameter**: Function or method parameter + - Supports type annotations + - Handles default values + - Supports generic types + +- **Decorator**: Python decorator for functions/classes + - Supports decorator arguments + - Handles import management + +- **Import**: Import statement management + - Supports from-import statements + - Handles multiple imports + - Manages import aliases + +- **ClassReference**: Reference to a class + - Used for imports and type hints + - Supports generic types + - Handles module paths + +- **CodeBlock**: Generic container for unsupported features + - Allows raw Python code + - Manages dependencies through references + - Preserves formatting + +## Publishing + +## Publish to npm + +In order to publish to npm `@amplication/python-ast` : + +1. Make sure to update the version in the package.json. +2. Run the following: + + +```sh +# From the monorepo root folder +npm i + +npx nx build python-ast + +cd ./dist/libs/python-ast + +``` + +To publish the package as "beta" run: + +``` +npm publish --access public --tag beta +``` + +To publish the package as "latest" run: + +```sh + +npm publish --access public + +``` + +## License + +MIT \ No newline at end of file diff --git a/libs/python-ast/jest.config.ts b/libs/python-ast/jest.config.ts new file mode 100644 index 0000000..8aa39ef --- /dev/null +++ b/libs/python-ast/jest.config.ts @@ -0,0 +1,11 @@ +/* eslint-disable */ +export default { + displayName: "python-ast", + preset: "../../jest.preset.js", + testEnvironment: "node", + transform: { + "^.+\\.[tj]s$": ["ts-jest", { tsconfig: "/tsconfig.spec.json" }], + }, + moduleFileExtensions: ["ts", "js", "html"], + coverageDirectory: "../../coverage/libs/python-ast", +}; diff --git a/libs/python-ast/package.json b/libs/python-ast/package.json new file mode 100644 index 0000000..7ce7628 --- /dev/null +++ b/libs/python-ast/package.json @@ -0,0 +1,11 @@ +{ + "name": "@amplication/python-ast", + "version": "0.1.0", + "description": "Python AST library in TypeScript", + "publishConfig": { + "access": "public" + }, + "dependencies": { + "@amplication/ast-types": "*" + } +} diff --git a/libs/python-ast/project.json b/libs/python-ast/project.json new file mode 100644 index 0000000..14728bc --- /dev/null +++ b/libs/python-ast/project.json @@ -0,0 +1,45 @@ +{ + "name": "python-ast", + "$schema": "../../node_modules/nx/schemas/project-schema.json", + "sourceRoot": "libs/python-ast/src", + "projectType": "library", + "targets": { + "build": { + "executor": "@nx/js:tsc", + "outputs": ["{options.outputPath}"], + "options": { + "outputPath": "dist/libs/python-ast", + "main": "libs/python-ast/src/index.ts", + "tsConfig": "libs/python-ast/tsconfig.lib.json", + "assets": ["libs/python-ast/*.md"] + } + }, + "lint": { + "executor": "@nx/linter:eslint", + "outputs": ["{options.outputFile}"], + "options": { + "lintFilePatterns": [ + "libs/python-ast/**/*.ts", + "libs/python-ast/package.json" + ] + } + }, + "test": { + "executor": "@nx/jest:jest", + "outputs": ["{workspaceRoot}/coverage/libs/python-ast"], + "options": { + "jestConfig": "libs/python-ast/jest.config.ts", + "passWithNoTests": true + } + }, + "update-snapshot": { + "executor": "@nx/jest:jest", + "outputs": ["{workspaceRoot}/coverage/libs/python-ast"], + "options": { + "jestConfig": "libs/python-ast/jest.config.ts", + "updateSnapshot": true + } + } + }, + "tags": [] +} diff --git a/libs/python-ast/src/index.ts b/libs/python-ast/src/index.ts new file mode 100644 index 0000000..a53bf7c --- /dev/null +++ b/libs/python-ast/src/index.ts @@ -0,0 +1,2 @@ +export * from "./lib/ast"; +export * from "./lib/core"; diff --git a/libs/python-ast/src/lib/ast/Assign.spec.ts b/libs/python-ast/src/lib/ast/Assign.spec.ts new file mode 100644 index 0000000..19d1652 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Assign.spec.ts @@ -0,0 +1,87 @@ +import { Assign } from "./Assign"; +import { ClassReference } from "./ClassReference"; + +describe("Assign", () => { + it("should generate a simple assignment", () => { + const assign = new Assign({ + target: "x", + value: "42", + }); + + expect(assign.toString()).toBe("x = 42"); + }); + + it("should generate an assignment with type annotation", () => { + const assign = new Assign({ + target: "name", + value: "'John'", + type: new ClassReference({ name: "str" }), + }); + + expect(assign.toString()).toBe("name: str = 'John'"); + }); + + it("should generate a class variable assignment", () => { + const assign = new Assign({ + target: "DEFAULT_TIMEOUT", + value: "30", + isClassVariable: true, + }); + + expect(assign.toString()).toBe("DEFAULT_TIMEOUT = 30"); + }); + + it("should generate an assignment with imported type", () => { + const assign = new Assign({ + target: "path", + value: "Path('/tmp')", + type: new ClassReference({ name: "Path", moduleName: "pathlib" }), + }); + + const output = assign.toString(); + expect(output).toContain("from pathlib import Path"); + expect(output).toContain("path: Path = Path('/tmp')"); + }); + + it("should generate an assignment with expression", () => { + const assign = new Assign({ + target: "total", + value: "x + y", + type: new ClassReference({ name: "float" }), + }); + + expect(assign.toString()).toBe("total: float = x + y"); + }); + + it("should generate an assignment with function call", () => { + const assign = new Assign({ + target: "result", + value: "calculate_total(items)", + type: new ClassReference({ name: "dict" }), + }); + + expect(assign.toString()).toBe("result: dict = calculate_total(items)"); + }); + + it("should generate an assignment with list comprehension", () => { + const assign = new Assign({ + target: "doubled", + value: "[x * 2 for x in numbers]", + type: new ClassReference({ name: "list" }), + }); + + expect(assign.toString()).toBe("doubled: list = [x * 2 for x in numbers]"); + }); + + it("should generate an assignment with complex type", () => { + const assign = new Assign({ + target: "items", + value: "[]", + type: new ClassReference({ name: "List", moduleName: "typing" }), + }); + + const output = assign.toString(); + expect(output).toContain("from typing import List"); + expect(output).toContain("items: List = []"); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Assign.ts b/libs/python-ast/src/lib/ast/Assign.ts new file mode 100644 index 0000000..66c8bb5 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Assign.ts @@ -0,0 +1,65 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { ClassReference } from "./ClassReference"; + +/** + * Configuration arguments for creating a Python assignment statement. + */ +export interface AssignArgs { + /** The target variable name */ + target: string; + /** The value to assign */ + value: string; + /** Type annotation for the target */ + type?: ClassReference; + /** Whether this is a class variable (vs. instance variable) */ + isClassVariable?: boolean; +} + +/** + * Represents a Python assignment statement. + * This class handles the generation of Python assignment statements, which can + * include type annotations and handle class vs. instance variables. + * + * @extends {AstNode} + */ +export class Assign extends AstNode { + /** The target variable name */ + public readonly target: string; + /** The value to assign */ + public readonly value: string; + /** Type annotation for the target */ + public readonly type?: ClassReference; + /** Whether this is a class variable */ + public readonly isClassVariable: boolean; + + /** + * Creates a new Python assignment statement. + * @param {AssignArgs} args - The configuration arguments + */ + constructor({ target, value, type, isClassVariable = false }: AssignArgs) { + super(); + this.target = target; + this.value = value; + this.type = type; + this.isClassVariable = isClassVariable; + } + + /** + * Writes the assignment statement to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + writer.write(this.target); + + if (this.type) { + writer.write(": "); + this.type.write(writer); + writer.write(" = "); + } else { + writer.write(" = "); + } + + writer.write(this.value); + } +} diff --git a/libs/python-ast/src/lib/ast/ClassDef.spec.ts b/libs/python-ast/src/lib/ast/ClassDef.spec.ts new file mode 100644 index 0000000..03eb815 --- /dev/null +++ b/libs/python-ast/src/lib/ast/ClassDef.spec.ts @@ -0,0 +1,159 @@ +import { ClassDef } from "./ClassDef"; +import { FunctionDef } from "./FunctionDef"; +import { Parameter } from "./Parameter"; +import { ClassReference } from "./ClassReference"; +import { CodeBlock } from "./CodeBlock"; +import { Decorator } from "./Decorator"; +import { Assign } from "./Assign"; + +describe("ClassDef", () => { + it("should generate a simple class with docstring", () => { + const classDef = new ClassDef({ + name: "Person", + docstring: "A class representing a person", + }); + + expect(classDef.toString()).toMatchSnapshot(); + expect(classDef.toString()).toContain("class Person:"); + expect(classDef.toString()).toContain( + '"""A class representing a person"""', + ); + }); + + it("should generate a class with inheritance", () => { + const classDef = new ClassDef({ + name: "Student", + bases: [new ClassReference({ name: "Person", moduleName: "models" })], + }); + + expect(classDef.toString()).toMatchSnapshot(); + expect(classDef.toString()).toContain("from models import Person"); + expect(classDef.toString()).toContain("class Student(Person):"); + }); + + it("should generate a class with methods", () => { + const classDef = new ClassDef({ + name: "Calculator", + }); + + const addMethod = new FunctionDef({ + name: "add", + parameters: [ + new Parameter({ name: "self" }), + new Parameter({ name: "a", type: new ClassReference({ name: "int" }) }), + new Parameter({ name: "b", type: new ClassReference({ name: "int" }) }), + ], + returnType: new ClassReference({ name: "int" }), + }); + + addMethod.addStatement(new CodeBlock({ code: "return a + b" })); + classDef.addMethod(addMethod); + + expect(classDef.toString()).toMatchSnapshot(); + expect(classDef.toString()).toContain( + "def add(self, a: int, b: int) -> int:", + ); + expect(classDef.toString()).toContain("return a + b"); + }); + + it("should generate a class with attributes", () => { + const classDef = new ClassDef({ + name: "Configuration", + }); + + classDef.addAttribute( + new Assign({ + target: "DEFAULT_TIMEOUT", + value: "30", + isClassVariable: true, + }), + ); + + classDef.addAttribute( + new Assign({ + target: "DEBUG", + value: "False", + isClassVariable: true, + }), + ); + + expect(classDef.toString()).toMatchSnapshot(); + expect(classDef.toString()).toContain("DEFAULT_TIMEOUT = 30"); + expect(classDef.toString()).toContain("DEBUG = False"); + }); + + it("should generate a class with decorators", () => { + const classDef = new ClassDef({ + name: "Singleton", + decorators: [ + new Decorator({ name: "dataclass", moduleName: "dataclasses" }), + ], + }); + + expect(classDef.toString()).toMatchSnapshot(); + expect(classDef.toString()).toContain("from dataclasses import dataclass"); + expect(classDef.toString()).toContain("@dataclass"); + }); + + it("should generate a complete class with multiple features", () => { + const classDef = new ClassDef({ + name: "User", + docstring: "A class representing a user in the system", + bases: [new ClassReference({ name: "BaseModel", moduleName: "models" })], + decorators: [ + new Decorator({ name: "dataclass", moduleName: "dataclasses" }), + ], + }); + + // Add class attributes + classDef.addAttribute( + new Assign({ + target: "table_name", + value: "'users'", + isClassVariable: true, + }), + ); + + // Add __init__ method + const initMethod = new FunctionDef({ + name: "__init__", + parameters: [ + new Parameter({ name: "self" }), + new Parameter({ + name: "username", + type: new ClassReference({ name: "str" }), + }), + new Parameter({ + name: "email", + type: new ClassReference({ name: "str" }), + }), + new Parameter({ + name: "is_active", + type: new ClassReference({ name: "bool" }), + default_: "True", + }), + ], + }); + initMethod.addStatement( + new CodeBlock({ + code: "self.username = username\nself.email = email\nself.is_active = is_active", + }), + ); + classDef.addMethod(initMethod); + + // Add toString method + const strMethod = new FunctionDef({ + name: "__str__", + parameters: [new Parameter({ name: "self" })], + returnType: new ClassReference({ name: "str" }), + }); + strMethod.addStatement( + new CodeBlock({ + code: "return f'User({self.username}, {self.email})'", + }), + ); + classDef.addMethod(strMethod); + + expect(classDef.toString()).toMatchSnapshot(); + }); +}); diff --git a/libs/python-ast/src/lib/ast/ClassDef.ts b/libs/python-ast/src/lib/ast/ClassDef.ts new file mode 100644 index 0000000..ca2b92f --- /dev/null +++ b/libs/python-ast/src/lib/ast/ClassDef.ts @@ -0,0 +1,220 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { Assign } from "./Assign"; +import { ClassReference } from "./ClassReference"; +import { CodeBlock } from "./CodeBlock"; +import { Decorator } from "./Decorator"; +import { FunctionDef } from "./FunctionDef"; + +/** + * Configuration arguments for creating a Python class. + */ +export interface ClassDefArgs { + /** The name of the class */ + name: string; + /** The module name containing the class */ + moduleName?: string; + /** The parent classes to inherit from */ + bases?: ClassReference[]; + /** Decorators to apply to the class */ + decorators?: Decorator[]; + /** Documentation string for the class */ + docstring?: string; +} + +/** + * Represents a Python class definition in the AST. + * This class handles the generation of Python class declarations including methods, + * attributes, and nested classes. + * + * @extends {AstNode} + */ +export class ClassDef extends AstNode { + /** The name of the class */ + public readonly name: string; + /** The module name containing the class */ + public readonly moduleName?: string; + /** The parent classes to inherit from */ + public readonly bases: ClassReference[]; + /** The decorators applied to the class */ + public readonly decorators: Decorator[]; + /** The documentation string for the class */ + public readonly docstring?: string; + /** The class reference for this class */ + public readonly reference: ClassReference; + + /** The methods defined in this class */ + private methods: FunctionDef[] = []; + /** The attributes defined in this class */ + private attributes: Assign[] = []; + /** The nested classes defined in this class */ + private nestedClasses: ClassDef[] = []; + /** Custom code blocks for unsupported features */ + private codeBlocks: CodeBlock[] = []; + + /** + * Creates a new Python class definition. + * @param {ClassDefArgs} args - The configuration arguments + */ + constructor({ + name, + moduleName, + bases = [], + decorators = [], + docstring, + }: ClassDefArgs) { + super(); + this.name = name; + this.moduleName = moduleName; + this.bases = bases; + this.decorators = decorators; + this.docstring = docstring; + + this.reference = new ClassReference({ + name: this.name, + moduleName: this.moduleName, + }); + } + + /** + * Adds a method to the class. + * @param {FunctionDef} method - The method to add + */ + public addMethod(method: FunctionDef): void { + this.methods.push(method); + } + + /** + * Adds an attribute to the class. + * @param {Assign} attribute - The attribute to add + */ + public addAttribute(attribute: Assign): void { + this.attributes.push(attribute); + } + + /** + * Adds a nested class to this class. + * @param {ClassDef} nestedClass - The nested class to add + */ + public addNestedClass(nestedClass: ClassDef): void { + this.nestedClasses.push(nestedClass); + } + + /** + * Adds a code block for unsupported features. + * @param {CodeBlock} codeBlock - The code block to add + */ + public addCodeBlock(codeBlock: CodeBlock): void { + this.codeBlocks.push(codeBlock); + } + + /** + * Writes the class definition and its members to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + // Write decorators + this.decorators.forEach((decorator) => { + decorator.write(writer); + writer.newLine(); + }); + + // Write class definition + writer.write(`class ${this.name}`); + + // Write base classes if any + if (this.bases.length > 0) { + writer.write("("); + this.bases.forEach((base, index) => { + base.write(writer); + if (index < this.bases.length - 1) { + writer.write(", "); + } + }); + writer.write(")"); + } + + writer.write(":"); + writer.newLine(); + writer.indent(); + + // Write docstring if provided + if (this.docstring) { + writer.writeLine(`"""${this.docstring}"""`); + if ( + this.attributes.length > 0 || + this.methods.length > 0 || + this.nestedClasses.length > 0 || + this.codeBlocks.length > 0 + ) { + writer.newLine(); + } + } + + // Write attributes + if (this.attributes.length > 0) { + this.attributes.forEach((attribute) => { + attribute.write(writer); + writer.newLine(); + }); + if ( + this.methods.length > 0 || + this.nestedClasses.length > 0 || + this.codeBlocks.length > 0 + ) { + writer.newLine(); + } + } + + // Write methods + this.methods.forEach((method, index, array) => { + method.write(writer); + if (index < array.length - 1) { + writer.newLine(); + } + }); + + if ( + this.methods.length > 0 && + (this.nestedClasses.length > 0 || this.codeBlocks.length > 0) + ) { + writer.newLine(); + writer.newLine(); + } + + // Write nested classes + this.nestedClasses.forEach((nestedClass, index, array) => { + nestedClass.write(writer); + if (index < array.length - 1) { + writer.newLine(); + writer.newLine(); + } + }); + + if (this.nestedClasses.length > 0 && this.codeBlocks.length > 0) { + writer.newLine(); + writer.newLine(); + } + + // Write additional code blocks + this.codeBlocks.forEach((codeBlock, index, array) => { + codeBlock.write(writer); + if (index < array.length - 1) { + writer.newLine(); + } + }); + + // If the class is empty, add a pass statement + if ( + this.attributes.length === 0 && + this.methods.length === 0 && + this.nestedClasses.length === 0 && + this.codeBlocks.length === 0 && + !this.docstring + ) { + writer.writeLine("pass"); + } + + writer.dedent(); + } +} diff --git a/libs/python-ast/src/lib/ast/ClassReference.spec.ts b/libs/python-ast/src/lib/ast/ClassReference.spec.ts new file mode 100644 index 0000000..43e2b63 --- /dev/null +++ b/libs/python-ast/src/lib/ast/ClassReference.spec.ts @@ -0,0 +1,109 @@ +import { ClassReference } from "./ClassReference"; +import { Writer } from "../core/Writer"; + +describe("ClassReference", () => { + it("should generate a simple class reference without module", () => { + const ref = new ClassReference({ + name: "str", + }); + + expect(ref.toString()).toBe("str"); + }); + + it("should generate a class reference with module", () => { + const ref = new ClassReference({ + name: "datetime", + moduleName: "datetime", + }); + + const output = ref.toString(); + expect(output).toContain("from datetime import datetime"); + expect(output).toContain("datetime"); + }); + + it("should add import to writer when module is specified", () => { + const ref = new ClassReference({ + name: "Path", + moduleName: "pathlib", + }); + const writer = new Writer({}); + + ref.write(writer); + + const output = writer.toString(); + expect(output).toContain("from pathlib import Path"); + expect(output).toContain("Path"); + }); + + it("should not add import to writer when module is not specified", () => { + const ref = new ClassReference({ + name: "int", + }); + const writer = new Writer({}); + + ref.write(writer); + + expect(writer.toString()).toBe("int"); + }); + + it("should handle multiple references to same class", () => { + const writer = new Writer({}); + const ref1 = new ClassReference({ + name: "UUID", + moduleName: "uuid", + }); + const ref2 = new ClassReference({ + name: "UUID", + moduleName: "uuid", + }); + + ref1.write(writer); + writer.write(", "); + ref2.write(writer); + + const output = writer.toString(); + expect(output).toContain("from uuid import UUID"); + expect(output.match(/from uuid import UUID/g)?.length).toBe(1); // Import should appear only once + expect(output).toContain("UUID, UUID"); + }); + + it("should handle nested module paths", () => { + const ref = new ClassReference({ + name: "HttpResponse", + moduleName: "django.http", + }); + + const output = ref.toString(); + expect(output).toContain("from django.http import HttpResponse"); + expect(output).toContain("HttpResponse"); + }); + + it("should handle type hints from typing module", () => { + const ref = new ClassReference({ + name: "List", + moduleName: "typing", + }); + + const output = ref.toString(); + expect(output).toContain("from typing import List"); + expect(output).toContain("List"); + }); + + it("should handle built-in types without imports", () => { + const builtinTypes = [ + "str", + "int", + "float", + "bool", + "list", + "dict", + "tuple", + "set", + ]; + + for (const type of builtinTypes) { + const ref = new ClassReference({ name: type }); + expect(ref.toString()).toBe(type); + } + }); +}); diff --git a/libs/python-ast/src/lib/ast/ClassReference.ts b/libs/python-ast/src/lib/ast/ClassReference.ts new file mode 100644 index 0000000..817b297 --- /dev/null +++ b/libs/python-ast/src/lib/ast/ClassReference.ts @@ -0,0 +1,64 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { Import } from "./Import"; + +/** + * Configuration arguments for creating a Python class reference. + */ +export interface ClassReferenceArgs { + /** The name of the class */ + name: string; + /** The module name containing the class */ + moduleName?: string; + /** Optional alias for the class when imported */ + alias?: string; +} + +/** + * Represents a reference to a Python class. + * This class handles the generation of class references in type annotations + * and manages the necessary imports for the referenced classes. + * + * @extends {AstNode} + */ +export class ClassReference extends AstNode { + /** The name of the class */ + public readonly name: string; + /** The module name containing the class */ + public readonly moduleName?: string; + /** Optional alias for the class */ + public readonly alias?: string; + /** The import node for this class reference */ + private readonly importNode: Import | null = null; + + /** + * Creates a new Python class reference. + * @param {ClassReferenceArgs} args - The configuration arguments + */ + constructor({ name, moduleName, alias }: ClassReferenceArgs) { + super(); + this.name = name; + this.moduleName = moduleName; + this.alias = alias; + + // Create an import node if the class is from a module + if (moduleName) { + this.importNode = new Import({ + moduleName, + names: [name], + alias, + }); + } + } + + /** + * Writes the class reference to the writer and adds it to the import collection. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + if (this.importNode) { + writer.addImport(this.importNode); + } + writer.write(this.alias || this.name); + } +} diff --git a/libs/python-ast/src/lib/ast/CodeBlock.spec.ts b/libs/python-ast/src/lib/ast/CodeBlock.spec.ts new file mode 100644 index 0000000..0ee64ae --- /dev/null +++ b/libs/python-ast/src/lib/ast/CodeBlock.spec.ts @@ -0,0 +1,69 @@ +import { CodeBlock } from "./CodeBlock"; +import { Writer } from "../core/Writer"; +import { Import } from "./Import"; + +describe("CodeBlock", () => { + it("should generate a simple code block", () => { + const block = new CodeBlock({ + code: 'print("Hello, World!")', + }); + + expect(block.toString()).toBe('print("Hello, World!")'); + }); + + it("should handle multiline code", () => { + const block = new CodeBlock({ + code: `x = 1 +y = 2 +print(x + y)`, + }); + + const output = block.toString(); + expect(output).toContain("x = 1"); + expect(output).toContain("y = 2"); + expect(output).toContain("print(x + y)"); + }); + + it("should handle code with imports", () => { + const block = new CodeBlock({ + code: "path = Path('/tmp')", + imports: [new Import({ moduleName: "pathlib", names: ["Path"] })], + }); + + const output = block.toString(); + expect(output).toContain("from pathlib import Path"); + expect(output).toContain("path = Path('/tmp')"); + }); + + it("should handle code with multiple imports", () => { + const block = new CodeBlock({ + code: `now = datetime.now() +uuid = UUID.uuid4()`, + imports: [ + new Import({ moduleName: "datetime", names: ["datetime"] }), + new Import({ moduleName: "uuid", names: ["UUID"] }), + ], + }); + + const output = block.toString(); + expect(output).toContain("from datetime import datetime"); + expect(output).toContain("from uuid import UUID"); + expect(output).toContain("now = datetime.now()"); + expect(output).toContain("uuid = UUID.uuid4()"); + }); + + it("should handle code generation function", () => { + const block = new CodeBlock({ + code: (writer: Writer) => { + writer.writeLine("def generate_id():"); + writer.indent(); + writer.writeLine("return UUID.uuid4()"); + writer.dedent(); + }, + }); + + const output = block.toString(); + expect(output).toContain("def generate_id():"); + expect(output).toContain(" return UUID.uuid4()"); + }); +}); diff --git a/libs/python-ast/src/lib/ast/CodeBlock.ts b/libs/python-ast/src/lib/ast/CodeBlock.ts new file mode 100644 index 0000000..56c340c --- /dev/null +++ b/libs/python-ast/src/lib/ast/CodeBlock.ts @@ -0,0 +1,61 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { Import } from "./Import"; + +/** + * Configuration arguments for creating a Python code block. + */ +export type CodeBlockArgs = + | { + /** The code to write */ + code: string; + /** A list of imports that are used in the code */ + imports?: Import[] | null; + } + | { + /** A function that writes code to the writer */ + code: (writer: Writer) => void; + imports?: never; + }; + +/** + * Represents a block of Python code in the AST. + * This class handles the generation of arbitrary Python code blocks, either as + * static strings or as dynamic code generation functions. It is used for language + * features not directly supported by specific AST node classes. + * + * @extends {AstNode} + */ +export class CodeBlock extends AstNode { + /** The code content or code generation function */ + private value: string | ((writer: Writer) => void); + /** The imports used in the code */ + private imports: Import[]; + + /** + * Creates a new Python code block. + * @param {CodeBlockArgs} args - The configuration arguments + */ + constructor(args: CodeBlockArgs) { + super(); + this.value = args.code; + this.imports = []; + + if ("imports" in args && args.imports) { + this.imports.push(...args.imports); + } + } + + /** + * Writes the code block to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + if (typeof this.value === "string") { + this.imports.forEach((import_) => writer.addImport(import_)); + writer.write(this.value); + } else { + this.value(writer); + } + } +} diff --git a/libs/python-ast/src/lib/ast/Decorator.spec.ts b/libs/python-ast/src/lib/ast/Decorator.spec.ts new file mode 100644 index 0000000..d8dd734 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Decorator.spec.ts @@ -0,0 +1,87 @@ +import { Decorator } from "./Decorator"; + +describe("Decorator", () => { + it("should generate a simple decorator", () => { + const decorator = new Decorator({ + name: "property", + }); + + expect(decorator.toString()).toBe("@property"); + }); + + it("should generate a decorator with arguments", () => { + const decorator = new Decorator({ + name: "route", + args: ["'/home'", "methods=['GET']"], + }); + + expect(decorator.toString()).toBe("@route('/home', methods=['GET'])"); + }); + + it("should generate a decorator with keyword arguments", () => { + const decorator = new Decorator({ + name: "validator", + kwargs: { + min_value: "0", + max_value: "100", + message: "'Value must be between 0 and 100'", + }, + }); + + expect(decorator.toString()).toBe( + "@validator(min_value=0, max_value=100, message='Value must be between 0 and 100')", + ); + }); + + it("should generate a decorator with both args and kwargs", () => { + const decorator = new Decorator({ + name: "click.option", + args: ["'--count'"], + kwargs: { + default: "1", + help: "'Number of iterations'", + }, + }); + + expect(decorator.toString()).toBe( + "@click.option('--count', default=1, help='Number of iterations')", + ); + }); + + it("should handle decorator with module import", () => { + const decorator = new Decorator({ + name: "dataclass", + moduleName: "dataclasses", + }); + + const output = decorator.toString(); + expect(output).toContain("from dataclasses import dataclass"); + expect(output).toContain("@dataclass"); + }); + + it("should handle decorator with module import and arguments", () => { + const decorator = new Decorator({ + name: "field", + moduleName: "dataclasses", + kwargs: { + default_factory: "list", + repr: "False", + }, + }); + + const output = decorator.toString(); + expect(output).toContain("from dataclasses import field"); + expect(output).toContain("@field(default_factory=list, repr=False)"); + }); + + it("should handle nested module paths", () => { + const decorator = new Decorator({ + name: "require_auth", + moduleName: "app.auth.decorators", + }); + + const output = decorator.toString(); + expect(output).toContain("from app.auth.decorators import require_auth"); + expect(output).toContain("@require_auth"); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Decorator.ts b/libs/python-ast/src/lib/ast/Decorator.ts new file mode 100644 index 0000000..53c44c0 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Decorator.ts @@ -0,0 +1,96 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { ClassReference } from "./ClassReference"; + +/** + * Configuration arguments for creating a Python decorator. + */ +export interface DecoratorArgs { + /** The name of the decorator */ + name: string; + /** The arguments to pass to the decorator */ + args?: string[]; + /** The keyword arguments to pass to the decorator */ + kwargs?: Record; + /** Optional module name for import */ + moduleName?: string; +} + +/** + * Represents a Python decorator. + * This class handles the generation of Python decorators, which are used to modify + * the behavior of functions, methods, and classes. + * + * @extends {AstNode} + */ +export class Decorator extends AstNode { + /** The name of the decorator */ + public readonly name: string; + /** The arguments to pass to the decorator */ + public readonly args: string[]; + /** The keyword arguments to pass to the decorator */ + public readonly kwargs: Record; + /** The module name containing the decorator */ + public readonly moduleName?: string; + /** The class reference for import handling */ + private readonly reference?: ClassReference; + + /** + * Creates a new Python decorator. + * @param {DecoratorArgs} args - The configuration arguments + */ + constructor({ name, args = [], kwargs = {}, moduleName }: DecoratorArgs) { + super(); + this.name = name; + this.args = args; + this.kwargs = kwargs; + this.moduleName = moduleName; + + if (moduleName) { + this.reference = new ClassReference({ + name: this.name, + moduleName: this.moduleName, + }); + } + } + + /** + * Writes the decorator to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + writer.write("@"); + + if (this.reference) { + this.reference.write(writer); + } else { + writer.write(this.name); + } + + // Add arguments and keyword arguments if provided + const hasArgs = this.args.length > 0; + const hasKwargs = Object.keys(this.kwargs).length > 0; + + if (hasArgs || hasKwargs) { + writer.write("("); + + // Write positional arguments + this.args.forEach((arg, index) => { + writer.write(arg); + if (index < this.args.length - 1 || hasKwargs) { + writer.write(", "); + } + }); + + // Write keyword arguments + Object.entries(this.kwargs).forEach(([key, value], index, array) => { + writer.write(`${key}=${value}`); + if (index < array.length - 1) { + writer.write(", "); + } + }); + + writer.write(")"); + } + } +} diff --git a/libs/python-ast/src/lib/ast/FunctionDef.spec.ts b/libs/python-ast/src/lib/ast/FunctionDef.spec.ts new file mode 100644 index 0000000..1af0a6a --- /dev/null +++ b/libs/python-ast/src/lib/ast/FunctionDef.spec.ts @@ -0,0 +1,199 @@ +import { FunctionDef } from "./FunctionDef"; +import { Parameter } from "./Parameter"; +import { ClassReference } from "./ClassReference"; +import { CodeBlock } from "./CodeBlock"; +import { Return } from "./Return"; +import { Decorator } from "./Decorator"; + +describe("FunctionDef", () => { + it("should generate a simple function", () => { + const func = new FunctionDef({ + name: "greet", + }); + + expect(func.toString()).toContain("def greet():"); + expect(func.toString()).toContain("pass"); + }); + + it("should generate a function with parameters", () => { + const func = new FunctionDef({ + name: "add", + parameters: [ + new Parameter({ + name: "a", + type: new ClassReference({ name: "int" }), + }), + new Parameter({ + name: "b", + type: new ClassReference({ name: "int" }), + }), + ], + returnType: new ClassReference({ name: "int" }), + }); + + func.addStatement(new Return({ value: "a + b" })); + + expect(func.toString()).toContain("def add(a: int, b: int) -> int:"); + expect(func.toString()).toContain("return a + b"); + }); + + it("should generate a function with docstring", () => { + const func = new FunctionDef({ + name: "calculate_total", + docstring: "Calculate the total of all items in the list.", + parameters: [ + new Parameter({ + name: "items", + type: new ClassReference({ name: "list" }), + }), + ], + }); + + expect(func.toString()).toContain("def calculate_total(items: list):"); + expect(func.toString()).toContain( + '"""Calculate the total of all items in the list."""', + ); + }); + + it("should generate a static method", () => { + const func = new FunctionDef({ + name: "create_from_dict", + isStatic: true, + parameters: [ + new Parameter({ + name: "data", + type: new ClassReference({ name: "dict" }), + }), + ], + }); + + expect(func.toString()).toContain("@staticmethod"); + expect(func.toString()).toContain("def create_from_dict(data: dict):"); + }); + + it("should generate a class method", () => { + const func = new FunctionDef({ + name: "from_json", + isClassMethod: true, + parameters: [ + new Parameter({ name: "cls" }), + new Parameter({ + name: "json_str", + type: new ClassReference({ name: "str" }), + }), + ], + }); + + expect(func.toString()).toContain("@classmethod"); + expect(func.toString()).toContain("def from_json(cls, json_str: str):"); + }); + + it("should generate an async function", () => { + const func = new FunctionDef({ + name: "fetch_data", + isAsync: true, + parameters: [ + new Parameter({ + name: "url", + type: new ClassReference({ name: "str" }), + }), + ], + }); + + expect(func.toString()).toContain("async def fetch_data(url: str):"); + }); + + it("should generate a function with decorators", () => { + const func = new FunctionDef({ + name: "get_name", + decorators: [ + new Decorator({ name: "property" }), + new Decorator({ + name: "validator", + args: ["'name'"], + moduleName: "validators", + }), + ], + }); + + const output = func.toString(); + expect(output).toContain("from validators import validator"); + expect(output).toContain("@property"); + expect(output).toContain("@validator('name')"); + expect(output).toContain("def get_name():"); + }); + + it("should generate a function with complex parameters", () => { + const func = new FunctionDef({ + name: "process_request", + parameters: [ + new Parameter({ + name: "request", + type: new ClassReference({ + name: "HttpRequest", + moduleName: "django.http", + }), + }), + new Parameter({ + name: "timeout", + type: new ClassReference({ name: "int" }), + default_: "30", + isKeywordOnly: true, + }), + new Parameter({ + name: "args", + type: new ClassReference({ name: "tuple" }), + isVariablePositional: true, + }), + new Parameter({ + name: "kwargs", + type: new ClassReference({ name: "dict" }), + isVariableKeyword: true, + }), + ], + returnType: new ClassReference({ + name: "HttpResponse", + moduleName: "django.http", + }), + }); + + const output = func.toString(); + expect(output).toContain( + "from django.http import HttpRequest, HttpResponse", + ); + expect(output).toContain( + "def process_request(request: HttpRequest, timeout: int = 30, *args: tuple, **kwargs: dict) -> HttpResponse:", + ); + }); + + it("should generate a function with multiple statements", () => { + const func = new FunctionDef({ + name: "process_data", + parameters: [ + new Parameter({ + name: "data", + type: new ClassReference({ name: "list" }), + }), + ], + }); + + func.addStatement( + new CodeBlock({ + code: "result = []", + }), + ); + func.addStatement( + new CodeBlock({ + code: "for item in data:\n result.append(item * 2)", + }), + ); + func.addStatement(new Return({ value: "result" })); + + const output = func.toString(); + expect(output).toContain("def process_data(data: list):"); + expect(output).toContain("result = []"); + expect(output).toContain("for item in data:"); + expect(output).toContain(" result.append(item * 2)"); + expect(output).toContain("return result"); + }); +}); diff --git a/libs/python-ast/src/lib/ast/FunctionDef.ts b/libs/python-ast/src/lib/ast/FunctionDef.ts new file mode 100644 index 0000000..eee6fcd --- /dev/null +++ b/libs/python-ast/src/lib/ast/FunctionDef.ts @@ -0,0 +1,164 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { ClassReference } from "./ClassReference"; +import { CodeBlock } from "./CodeBlock"; +import { Decorator } from "./Decorator"; +import { Parameter } from "./Parameter"; +import { Return } from "./Return"; + +/** + * Configuration arguments for creating a Python function. + */ +export interface FunctionDefArgs { + /** The name of the function */ + name: string; + /** The parameters of the function */ + parameters?: Parameter[]; + /** Decorators to apply to the function */ + decorators?: Decorator[]; + /** Documentation string for the function */ + docstring?: string; + /** Return type annotation */ + returnType?: ClassReference; + /** Whether this is a static method */ + isStatic?: boolean; + /** Whether this is a class method */ + isClassMethod?: boolean; + /** Whether this is an async function */ + isAsync?: boolean; +} + +/** + * Represents a Python function definition in the AST. + * This class handles the generation of Python function declarations including + * parameters, decorators, and return type annotations. + * + * @extends {AstNode} + */ +export class FunctionDef extends AstNode { + /** The name of the function */ + public readonly name: string; + /** The parameters of the function */ + public readonly parameters: Parameter[]; + /** Decorators applied to the function */ + public readonly decorators: Decorator[]; + /** The documentation string */ + public readonly docstring?: string; + /** The return type annotation */ + public readonly returnType?: ClassReference; + /** Whether this is a static method */ + public readonly isStatic: boolean; + /** Whether this is a class method */ + public readonly isClassMethod: boolean; + /** Whether this is an async function */ + public readonly isAsync: boolean; + + /** The body of the function */ + private body: (CodeBlock | Return)[] = []; + + /** + * Creates a new Python function definition. + * @param {FunctionDefArgs} args - The configuration arguments + */ + constructor({ + name, + parameters = [], + decorators = [], + docstring, + returnType, + isStatic = false, + isClassMethod = false, + isAsync = false, + }: FunctionDefArgs) { + super(); + this.name = name; + this.parameters = parameters; + this.decorators = [...decorators]; + this.docstring = docstring; + this.returnType = returnType; + this.isStatic = isStatic; + this.isClassMethod = isClassMethod; + this.isAsync = isAsync; + + // Add appropriate decorators for static and class methods + if ( + this.isStatic && + !this.decorators.some((d) => d.name === "staticmethod") + ) { + this.decorators.push(new Decorator({ name: "staticmethod" })); + } + if ( + this.isClassMethod && + !this.decorators.some((d) => d.name === "classmethod") + ) { + this.decorators.push(new Decorator({ name: "classmethod" })); + } + } + + /** + * Adds a statement to the function body. + * @param {CodeBlock | Return} statement - The statement to add + */ + public addStatement(statement: CodeBlock | Return): void { + this.body.push(statement); + } + + /** + * Writes the function definition and its body to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + // Write decorators + this.decorators.forEach((decorator) => { + decorator.write(writer); + writer.newLine(); + }); + + // Write function definition + if (this.isAsync) { + writer.write("async "); + } + writer.write(`def ${this.name}(`); + + // Write parameters + this.parameters.forEach((param, index) => { + param.write(writer); + if (index < this.parameters.length - 1) { + writer.write(", "); + } + }); + + writer.write(")"); + + // Write return type annotation if provided + if (this.returnType) { + writer.write(" -> "); + this.returnType.write(writer); + } + + writer.write(":"); + writer.newLine(); + writer.indent(); + + // Write docstring if provided + if (this.docstring) { + writer.writeLine(`"""${this.docstring}"""`); + if (this.body.length > 0) { + writer.newLine(); + } + } + + // Write function body + if (this.body.length > 0) { + this.body.forEach((statement) => { + statement.write(writer); + writer.newLine(); + }); + } else { + // Empty function body needs a pass statement + writer.writeLine("pass"); + } + + writer.dedent(); + } +} diff --git a/libs/python-ast/src/lib/ast/Import.spec.ts b/libs/python-ast/src/lib/ast/Import.spec.ts new file mode 100644 index 0000000..1a81a7e --- /dev/null +++ b/libs/python-ast/src/lib/ast/Import.spec.ts @@ -0,0 +1,75 @@ +import { Import } from "./Import"; + +describe("Import", () => { + it("should generate a simple import", () => { + const importStmt = new Import({ + moduleName: "os", + }); + + expect(importStmt.toString()).toContain("import os"); + }); + + it("should generate an import with alias", () => { + const importStmt = new Import({ + moduleName: "numpy", + alias: "np", + }); + + expect(importStmt.toString()).toContain("import numpy as np"); + }); + + it("should generate a from import", () => { + const importStmt = new Import({ + moduleName: "datetime", + names: ["datetime", "timedelta"], + }); + + expect(importStmt.toString()).toContain( + "from datetime import datetime, timedelta", + ); + }); + + it("should generate a from import with single name and alias", () => { + const importStmt = new Import({ + moduleName: "pandas", + names: ["DataFrame"], + alias: "df", + }); + + expect(importStmt.toString()).toContain( + "from pandas import DataFrame as df", + ); + }); + + it("should handle nested module paths", () => { + const importStmt = new Import({ + moduleName: "django.http", + names: ["HttpResponse", "JsonResponse"], + }); + + expect(importStmt.toString()).toContain( + "from django.http import HttpResponse, JsonResponse", + ); + }); + + it("should handle empty names array as simple import", () => { + const importStmt = new Import({ + moduleName: "sys", + names: [], + }); + + expect(importStmt.toString()).toContain("import sys"); + }); + + it("should ignore alias when multiple names are imported", () => { + const importStmt = new Import({ + moduleName: "os.path", + names: ["join", "dirname", "basename"], + alias: "ignored", + }); + + expect(importStmt.toString()).toContain( + `from os.path import basename, dirname, join`, + ); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Import.ts b/libs/python-ast/src/lib/ast/Import.ts new file mode 100644 index 0000000..b36fc15 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Import.ts @@ -0,0 +1,49 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; + +/** + * Configuration arguments for creating a Python import statement. + */ +export interface ImportArgs { + /** The module name to import */ + moduleName: string; + /** The names to import from the module (if using 'from ... import ...') */ + names?: string[]; + /** The alias for the module or imported name (if using 'import ... as ...') */ + alias?: string; +} + +/** + * Represents a Python import statement. + * This class handles the generation of Python import statements, including + * simple imports, from imports, and imports with aliases. + * + * @extends {AstNode} + */ +export class Import extends AstNode { + /** The module name to import */ + public readonly moduleName: string; + /** The names to import from the module */ + public readonly names: string[]; + /** The alias for the module or imported name */ + public readonly alias?: string; + + /** + * Creates a new Python import statement. + * @param {ImportArgs} args - The configuration arguments + */ + constructor({ moduleName, names = [], alias }: ImportArgs) { + super(); + this.moduleName = moduleName; + this.names = names; + this.alias = alias; + } + + /** + * Writes the import statement to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + writer.addImport(this); + } +} diff --git a/libs/python-ast/src/lib/ast/Module.spec.ts b/libs/python-ast/src/lib/ast/Module.spec.ts new file mode 100644 index 0000000..1c24211 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Module.spec.ts @@ -0,0 +1,222 @@ +import { Module } from "./Module"; +import { ClassDef } from "./ClassDef"; +import { FunctionDef } from "./FunctionDef"; +import { Parameter } from "./Parameter"; +import { ClassReference } from "./ClassReference"; +import { CodeBlock } from "./CodeBlock"; +import { Import } from "./Import"; + +describe("Module", () => { + it("should generate an empty module with docstring", () => { + const module = new Module({ + name: "config", + docstring: "Configuration module for the application", + }); + + expect(module.toString()).toMatchSnapshot(); + expect(module.toString()).toContain( + '"""Configuration module for the application"""', + ); + }); + + it("should generate a module with imports", () => { + const module = new Module({ + name: "utils", + }); + + module.addImport( + new Import({ + moduleName: "os", + }), + ); + + module.addImport( + new Import({ + moduleName: "datetime", + names: ["datetime", "timezone"], + }), + ); + + expect(module.toString()).toMatchSnapshot(); + expect(module.toString()).toContain("import os"); + expect(module.toString()).toContain( + "from datetime import datetime, timezone", + ); + }); + + it("should generate a module with functions", () => { + const module = new Module({ + name: "math_utils", + }); + + const squareFunction = new FunctionDef({ + name: "square", + parameters: [ + new Parameter({ + name: "x", + type: new ClassReference({ name: "float" }), + }), + ], + returnType: new ClassReference({ name: "float" }), + }); + squareFunction.addStatement(new CodeBlock({ code: "return x * x" })); + module.addFunction(squareFunction); + + const cubeFunction = new FunctionDef({ + name: "cube", + parameters: [ + new Parameter({ + name: "x", + type: new ClassReference({ name: "float" }), + }), + ], + returnType: new ClassReference({ name: "float" }), + }); + cubeFunction.addStatement(new CodeBlock({ code: "return x * x * x" })); + module.addFunction(cubeFunction); + + expect(module.toString()).toMatchSnapshot(); + expect(module.toString()).toContain("def square(x: float) -> float:"); + expect(module.toString()).toContain("return x * x"); + expect(module.toString()).toContain("def cube(x: float) -> float:"); + expect(module.toString()).toContain("return x * x * x"); + }); + + it("should generate a module with classes", () => { + const module = new Module({ + name: "models", + }); + + const userClass = new ClassDef({ + name: "User", + docstring: "User model", + }); + module.addClass(userClass); + + const productClass = new ClassDef({ + name: "Product", + docstring: "Product model", + }); + module.addClass(productClass); + + expect(module.toString()).toMatchSnapshot(); + expect(module.toString()).toContain("class User:"); + expect(module.toString()).toContain('"""User model"""'); + expect(module.toString()).toContain("class Product:"); + expect(module.toString()).toContain('"""Product model"""'); + }); + + it("should generate a module with global code", () => { + const module = new Module({ + name: "settings", + }); + + module.addCodeBlock( + new CodeBlock({ + code: 'DEBUG = True\nVERSION = "1.0.0"\nBASE_DIR = os.path.dirname(os.path.abspath(__file__))', + imports: [new Import({ moduleName: "os", names: ["os"] })], + }), + ); + + expect(module.toString()).toMatchSnapshot(); + expect(module.toString()).toContain("import os"); + expect(module.toString()).toContain("DEBUG = True"); + expect(module.toString()).toContain('VERSION = "1.0.0"'); + expect(module.toString()).toContain( + "BASE_DIR = os.path.dirname(os.path.abspath(__file__))", + ); + }); + + it("should generate a complete module with all features and maintain order", () => { + const module = new Module({ + name: "app", + docstring: "Main application module", + }); + + // Add imports + module.addImport( + new Import({ + moduleName: "os", + }), + ); + module.addImport( + new Import({ + moduleName: "logging", + }), + ); + + // Add global code + module.addCodeBlock( + new CodeBlock({ + code: 'logger = logging.getLogger(__name__)\nAPP_NAME = "MyApp"\nVERSION = "1.0.0"', + }), + ); + + // Add a utility function + const configFunction = new FunctionDef({ + name: "configure_app", + parameters: [ + new Parameter({ + name: "debug", + type: new ClassReference({ name: "bool" }), + default_: "False", + }), + ], + }); + configFunction.addStatement( + new CodeBlock({ + code: 'logger.info("Configuring app with debug=%s", debug)\nreturn {"debug": debug, "version": VERSION}', + }), + ); + module.addFunction(configFunction); + + // Add an application class + const appClass = new ClassDef({ + name: "Application", + docstring: "Main application class", + }); + + const initMethod = new FunctionDef({ + name: "__init__", + parameters: [ + new Parameter({ name: "self" }), + new Parameter({ + name: "config", + type: new ClassReference({ name: "dict" }), + }), + ], + }); + initMethod.addStatement( + new CodeBlock({ + code: 'self.config = config\nlogger.info("Application initialized")', + }), + ); + appClass.addMethod(initMethod); + + const runMethod = new FunctionDef({ + name: "run", + parameters: [new Parameter({ name: "self" })], + }); + runMethod.addStatement( + new CodeBlock({ + code: 'logger.info("Running application %s", APP_NAME)\nprint("Application running...")', + }), + ); + appClass.addMethod(runMethod); + + module.addClass(appClass); + + expect(module.toString()).toMatchSnapshot(); + + // Additional test to verify order + const output = module.toString(); + const codeBlockIndex = output.indexOf( + "logger = logging.getLogger(__name__)", + ); + const functionIndex = output.indexOf("def configure_app"); + const classIndex = output.indexOf("class Application:"); + + expect(codeBlockIndex).toBeLessThan(functionIndex); + expect(functionIndex).toBeLessThan(classIndex); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Module.ts b/libs/python-ast/src/lib/ast/Module.ts new file mode 100644 index 0000000..ba2b4b9 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Module.ts @@ -0,0 +1,106 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { ClassDef } from "./ClassDef"; +import { CodeBlock } from "./CodeBlock"; +import { FunctionDef } from "./FunctionDef"; +import { Import } from "./Import"; + +/** + * Configuration arguments for creating a Python module. + */ +export interface ModuleArgs { + /** The name of the module */ + name: string; + /** The docstring for the module */ + docstring?: string; +} + +/** + * Represents a Python module in the AST. + * This class is the top-level container for Python code, including imports, + * functions, classes, and global code. + * + * @extends {AstNode} + */ +export class Module extends AstNode { + /** The name of the module */ + public readonly name: string; + /** The docstring for the module */ + public readonly docstring?: string; + + /** The imports defined in this module */ + private imports: Import[] = []; + /** The children nodes (functions, classes, code blocks) in order of addition */ + private children: (FunctionDef | ClassDef | CodeBlock)[] = []; + + /** + * Creates a new Python module. + * @param {ModuleArgs} args - The configuration arguments + */ + constructor({ name, docstring }: ModuleArgs) { + super(); + this.name = name; + this.docstring = docstring; + } + + /** + * Adds an import statement to the module. + * @param {Import} importStatement - The import statement to add + */ + public addImport(importStatement: Import): void { + this.imports.push(importStatement); + } + + /** + * Adds a function to the module. + * @param {FunctionDef} functionDef - The function to add + */ + public addFunction(functionDef: FunctionDef): void { + this.children.push(functionDef); + } + + /** + * Adds a class to the module. + * @param {ClassDef} classDef - The class to add + */ + public addClass(classDef: ClassDef): void { + this.children.push(classDef); + } + + /** + * Adds a global code block to the module. + * @param {CodeBlock} codeBlock - The code block to add + */ + public addCodeBlock(codeBlock: CodeBlock): void { + this.children.push(codeBlock); + } + + /** + * Writes the module and its contents to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + // Write imports + if (this.imports.length > 0) { + this.imports.forEach((import_) => import_.write(writer)); + writer.newLine(); + } + + // Write module docstring if provided + if (this.docstring) { + writer.writeLine(`"""${this.docstring}"""`); + writer.newLine(); + } + + // Write all children in order + if (this.children.length > 0) { + this.children.forEach((child, index) => { + child.write(writer); + if (index < this.children.length - 1) { + writer.newLine(); + writer.newLine(); + } + }); + } + } +} diff --git a/libs/python-ast/src/lib/ast/Parameter.spec.ts b/libs/python-ast/src/lib/ast/Parameter.spec.ts new file mode 100644 index 0000000..350c9c7 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Parameter.spec.ts @@ -0,0 +1,147 @@ +import { Parameter } from "./Parameter"; +import { ClassReference } from "./ClassReference"; + +describe("Parameter", () => { + it("should generate a parameter with a name", () => { + const param = new Parameter({ name: "x" }); + expect(param.toString()).toBe("x"); + }); + + it("should generate a parameter with a type annotation", () => { + const param = new Parameter({ + name: "name", + type: new ClassReference({ name: "str" }), + }); + expect(param.toString()).toBe("name: str"); + }); + + it("should generate a parameter with a default value", () => { + const param = new Parameter({ + name: "enabled", + type: new ClassReference({ name: "bool" }), + default_: "False", + }); + expect(param.toString()).toBe("enabled: bool = False"); + }); + + it("should generate a parameter with a numeric type", () => { + const param = new Parameter({ + name: "count", + type: new ClassReference({ name: "int" }), + default_: "0", + }); + expect(param.toString()).toBe("count: int = 0"); + }); + + it("should generate a parameter with a float type", () => { + const param = new Parameter({ + name: "value", + type: new ClassReference({ name: "float" }), + default_: "0.0", + }); + expect(param.toString()).toBe("value: float = 0.0"); + }); + + it("should generate a parameter with a tuple type", () => { + const param = new Parameter({ + name: "coordinates", + type: new ClassReference({ name: "tuple" }), + default_: "(0, 0)", + }); + expect(param.toString()).toBe("coordinates: tuple = (0, 0)"); + }); + + it("should generate a parameter with a dict type", () => { + const param = new Parameter({ + name: "options", + type: new ClassReference({ name: "dict" }), + default_: "{}", + }); + expect(param.toString()).toBe("options: dict = {}"); + }); + + it("should generate a parameter with an imported type", () => { + const param = new Parameter({ + name: "path", + type: new ClassReference({ name: "Path", moduleName: "pathlib" }), + default_: 'Path(".")', + }); + expect(param.toString()).toContain("from pathlib import Path"); + expect(param.toString()).toContain('path: Path = Path(".")'); + }); + + it("should generate a keyword-only parameter", () => { + const param = new Parameter({ + name: "timeout", + type: new ClassReference({ name: "int" }), + default_: "30", + isKeywordOnly: true, + }); + expect(param.toString()).toBe("timeout: int = 30"); + }); + + it("should generate a positional-only parameter", () => { + const param = new Parameter({ + name: "x", + type: new ClassReference({ name: "float" }), + isPositionalOnly: true, + }); + expect(param.toString()).toBe("x: float"); + }); + + it("should generate a variable positional parameter", () => { + const param = new Parameter({ + name: "args", + type: new ClassReference({ name: "tuple" }), + isVariablePositional: true, + }); + expect(param.toString()).toBe("*args: tuple"); + }); + + it("should generate a variable keyword parameter", () => { + const param = new Parameter({ + name: "kwargs", + type: new ClassReference({ name: "dict" }), + isVariableKeyword: true, + }); + expect(param.toString()).toBe("**kwargs: dict"); + }); + + it("should generate a parameter with an imported type and no default", () => { + const param = new Parameter({ + name: "path", + type: new ClassReference({ name: "Path", moduleName: "pathlib" }), + }); + expect(param.toString()).toContain("from pathlib import Path"); + expect(param.toString()).toContain("path: Path"); + }); + + it("should handle parameter with imported type", () => { + const param = new Parameter({ + name: "path", + type: new ClassReference({ name: "Path", moduleName: "pathlib" }), + }); + + const output = param.toString(); + expect(output).toContain("from pathlib import Path"); + expect(output).toContain("path: Path"); + }); + + it("should throw error when multiple parameter types are specified", () => { + expect(() => { + new Parameter({ + name: "invalid", + isKeywordOnly: true, + isPositionalOnly: true, + }); + }).toThrow(); + + expect(() => { + new Parameter({ + name: "invalid", + isVariablePositional: true, + isVariableKeyword: true, + }); + }).toThrow(); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Parameter.ts b/libs/python-ast/src/lib/ast/Parameter.ts new file mode 100644 index 0000000..98c61d7 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Parameter.ts @@ -0,0 +1,108 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; +import { ClassReference } from "./ClassReference"; + +/** + * Configuration arguments for creating a Python parameter. + */ +export interface ParameterArgs { + /** The name of the parameter */ + name: string; + /** The type annotation for the parameter */ + type?: ClassReference; + /** The default value for the parameter */ + default_?: string; + /** Whether this is a keyword-only parameter (after *) */ + isKeywordOnly?: boolean; + /** Whether this is a positional-only parameter (before /) */ + isPositionalOnly?: boolean; + /** Whether this is a variable positional parameter (*args) */ + isVariablePositional?: boolean; + /** Whether this is a variable keyword parameter (**kwargs) */ + isVariableKeyword?: boolean; +} + +/** + * Represents a Python parameter in a function or method definition. + * This class handles the generation of parameter declarations including + * type annotations and default values. + * + * @extends {AstNode} + */ +export class Parameter extends AstNode { + /** The name of the parameter */ + public readonly name: string; + /** The type annotation for the parameter */ + public readonly type?: ClassReference; + /** The default value for the parameter */ + public readonly default_?: string; + /** Whether this is a keyword-only parameter */ + public readonly isKeywordOnly: boolean; + /** Whether this is a positional-only parameter */ + public readonly isPositionalOnly: boolean; + /** Whether this is a variable positional parameter */ + public readonly isVariablePositional: boolean; + /** Whether this is a variable keyword parameter */ + public readonly isVariableKeyword: boolean; + + /** + * Creates a new Python parameter. + * @param {ParameterArgs} args - The configuration arguments + */ + constructor({ + name, + type, + default_, + isKeywordOnly = false, + isPositionalOnly = false, + isVariablePositional = false, + isVariableKeyword = false, + }: ParameterArgs) { + super(); + this.name = name; + this.type = type; + this.default_ = default_; + this.isKeywordOnly = isKeywordOnly; + this.isPositionalOnly = isPositionalOnly; + this.isVariablePositional = isVariablePositional; + this.isVariableKeyword = isVariableKeyword; + + // Validate parameter configuration + if ( + [ + isKeywordOnly, + isPositionalOnly, + isVariablePositional, + isVariableKeyword, + ].filter(Boolean).length > 1 + ) { + throw new Error( + "A parameter can only be one of: keyword-only, positional-only, variable positional, or variable keyword", + ); + } + } + + /** + * Writes the parameter declaration to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + if (this.isVariablePositional) { + writer.write("*"); + } else if (this.isVariableKeyword) { + writer.write("**"); + } + + writer.write(this.name); + + if (this.type) { + writer.write(": "); + this.type.write(writer); + } + + if (this.default_) { + writer.write(" = "); + writer.write(this.default_); + } + } +} diff --git a/libs/python-ast/src/lib/ast/Return.spec.ts b/libs/python-ast/src/lib/ast/Return.spec.ts new file mode 100644 index 0000000..14e78d4 --- /dev/null +++ b/libs/python-ast/src/lib/ast/Return.spec.ts @@ -0,0 +1,57 @@ +import { Return } from "./Return"; + +describe("Return", () => { + it("should generate a simple return statement", () => { + const returnStmt = new Return(); + + expect(returnStmt.toString()).toBe("return"); + }); + + it("should generate a return statement with value", () => { + const returnStmt = new Return({ + value: "42", + }); + + expect(returnStmt.toString()).toBe("return 42"); + }); + + it("should generate a return statement with string value", () => { + const returnStmt = new Return({ + value: "'Hello, World!'", + }); + + expect(returnStmt.toString()).toBe("return 'Hello, World!'"); + }); + + it("should generate a return statement with expression", () => { + const returnStmt = new Return({ + value: "x + y", + }); + + expect(returnStmt.toString()).toBe("return x + y"); + }); + + it("should generate a return statement with function call", () => { + const returnStmt = new Return({ + value: "calculate_total(items)", + }); + + expect(returnStmt.toString()).toBe("return calculate_total(items)"); + }); + + it("should generate a return statement with dictionary", () => { + const returnStmt = new Return({ + value: "{'name': name, 'age': age}", + }); + + expect(returnStmt.toString()).toBe("return {'name': name, 'age': age}"); + }); + + it("should generate a return statement with list comprehension", () => { + const returnStmt = new Return({ + value: "[x * 2 for x in numbers]", + }); + + expect(returnStmt.toString()).toBe("return [x * 2 for x in numbers]"); + }); +}); diff --git a/libs/python-ast/src/lib/ast/Return.ts b/libs/python-ast/src/lib/ast/Return.ts new file mode 100644 index 0000000..b5f6dcd --- /dev/null +++ b/libs/python-ast/src/lib/ast/Return.ts @@ -0,0 +1,43 @@ +import { AstNode } from "../core/AstNode"; +import { Writer } from "../core/Writer"; + +/** + * Configuration arguments for creating a Python return statement. + */ +export interface ReturnArgs { + /** The value to return (if any) */ + value?: string; +} + +/** + * Represents a Python return statement. + * This class handles the generation of Python return statements, which can + * either return a value or return nothing. + * + * @extends {AstNode} + */ +export class Return extends AstNode { + /** The value to return */ + public readonly value?: string; + + /** + * Creates a new Python return statement. + * @param {ReturnArgs} args - The configuration arguments + */ + constructor({ value }: ReturnArgs = {}) { + super(); + this.value = value; + } + + /** + * Writes the return statement to the writer. + * @param {Writer} writer - The writer to write to + */ + public write(writer: Writer): void { + if (this.value) { + writer.write(`return ${this.value}`); + } else { + writer.write("return"); + } + } +} diff --git a/libs/python-ast/src/lib/ast/__snapshots__/ClassDef.spec.ts.snap b/libs/python-ast/src/lib/ast/__snapshots__/ClassDef.spec.ts.snap new file mode 100644 index 0000000..f46aefb --- /dev/null +++ b/libs/python-ast/src/lib/ast/__snapshots__/ClassDef.spec.ts.snap @@ -0,0 +1,58 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`ClassDef should generate a class with attributes 1`] = ` +"class Configuration: + DEFAULT_TIMEOUT = 30 + DEBUG = False +" +`; + +exports[`ClassDef should generate a class with decorators 1`] = ` +"from dataclasses import dataclass + +@dataclass +class Singleton: + pass +" +`; + +exports[`ClassDef should generate a class with inheritance 1`] = ` +"from models import Person + +class Student(Person): + pass +" +`; + +exports[`ClassDef should generate a class with methods 1`] = ` +"class Calculator: + def add(self, a: int, b: int) -> int: + return a + b +" +`; + +exports[`ClassDef should generate a complete class with multiple features 1`] = ` +"from dataclasses import dataclass +from models import BaseModel + +@dataclass +class User(BaseModel): + """A class representing a user in the system""" + + table_name = 'users' + + def __init__(self, username: str, email: str, is_active: bool = True): + self.username = username + self.email = email + self.is_active = is_active + + def __str__(self) -> str: + return f'User({self.username}, {self.email})' +" +`; + +exports[`ClassDef should generate a simple class with docstring 1`] = ` +"class Person: + """A class representing a person""" +" +`; diff --git a/libs/python-ast/src/lib/ast/__snapshots__/Module.spec.ts.snap b/libs/python-ast/src/lib/ast/__snapshots__/Module.spec.ts.snap new file mode 100644 index 0000000..2bf69e3 --- /dev/null +++ b/libs/python-ast/src/lib/ast/__snapshots__/Module.spec.ts.snap @@ -0,0 +1,72 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`Module should generate a complete module with all features and maintain order 1`] = ` +"import os +import logging + + +"""Main application module""" + +logger = logging.getLogger(__name__) +APP_NAME = "MyApp" +VERSION = "1.0.0" + +def configure_app(debug: bool = False): + logger.info("Configuring app with debug=%s", debug) + return {"debug": debug, "version": VERSION} + + +class Application: + """Main application class""" + + def __init__(self, config: dict): + self.config = config + logger.info("Application initialized") + + def run(self): + logger.info("Running application %s", APP_NAME) + print("Application running...") +" +`; + +exports[`Module should generate a module with classes 1`] = ` +"class User: + """User model""" + + +class Product: + """Product model""" +" +`; + +exports[`Module should generate a module with functions 1`] = ` +"def square(x: float) -> float: + return x * x + + +def cube(x: float) -> float: + return x * x * x +" +`; + +exports[`Module should generate a module with global code 1`] = ` +"from os import os + +DEBUG = True +VERSION = "1.0.0" +BASE_DIR = os.path.dirname(os.path.abspath(__file__))" +`; + +exports[`Module should generate a module with imports 1`] = ` +"import os +from datetime import datetime, timezone + + +" +`; + +exports[`Module should generate an empty module with docstring 1`] = ` +""""Configuration module for the application""" + +" +`; diff --git a/libs/python-ast/src/lib/ast/__snapshots__/integration.spec.ts.snap b/libs/python-ast/src/lib/ast/__snapshots__/integration.spec.ts.snap new file mode 100644 index 0000000..40cf8c7 --- /dev/null +++ b/libs/python-ast/src/lib/ast/__snapshots__/integration.spec.ts.snap @@ -0,0 +1,79 @@ +// Jest Snapshot v1, https://goo.gl/fbAQLP + +exports[`Python AST Integration Complex Class Generation should generate a valid Python class with methods and type hints 1`] = ` +"from base_module import BaseClass +from typing import Optional + +"""Example module demonstrating Python AST generation""" + +class Person(BaseClass): + """A class representing a person with name and age.""" + + default_age: int = 0 + + def __init__(self, name: str, age: Optional = None): + """Initialize a Person instance.""" + + self.name = name + self.age = age if age is not None else self.default_age + + @staticmethod + def from_dict(data: dict) -> Person: + """Create a Person instance from a dictionary.""" + + return Person(data["name"], data.get("age")) + + @property + def get_info(self) -> str: + return f"{self.name} ({self.age} years old)" +" +`; + +exports[`Python AST Integration should generate a complete, realistic Python file 1`] = ` +"from dataclasses import dataclass, field +from datetime import datetime +from typing import Dict, List, Optional +from uuid import UUID + + +"""Module containing the User data model and related functionality.""" + +@dataclass +class User: + """Represents a user in the system with their associated metadata.""" + + id: UUID = UUID.uuid4() + username: str = '' + email: str = '' + created_at: datetime = datetime.now() + roles: List = field(default_factory=list) + metadata: Dict = field(default_factory=dict) + + @classmethod + def create(cls, username: str, email: str, roles: List = None) -> User: + return cls( + id=UUID.uuid4(), + username=username, + email=email, + roles=roles or [], + created_at=datetime.now(), + metadata={}) + + def add_role(self, role: str) -> None: + if role not in self.roles: + self.roles.append(role) + + @property + def is_admin(self) -> bool: + return 'admin' in self.roles + + + + +# Example usage: +if __name__ == "__main__": + user = User.create("john_doe", "john@example.com") + user.add_role("admin") + print(f"Is admin: {user.is_admin}") +" +`; diff --git a/libs/python-ast/src/lib/ast/index.ts b/libs/python-ast/src/lib/ast/index.ts new file mode 100644 index 0000000..9587e20 --- /dev/null +++ b/libs/python-ast/src/lib/ast/index.ts @@ -0,0 +1,10 @@ +export * from "./Assign"; +export * from "./ClassDef"; +export * from "./ClassReference"; +export * from "./CodeBlock"; +export * from "./Decorator"; +export * from "./FunctionDef"; +export * from "./Import"; +export * from "./Module"; +export * from "./Parameter"; +export * from "./Return"; diff --git a/libs/python-ast/src/lib/ast/integration.spec.ts b/libs/python-ast/src/lib/ast/integration.spec.ts new file mode 100644 index 0000000..ff14839 --- /dev/null +++ b/libs/python-ast/src/lib/ast/integration.spec.ts @@ -0,0 +1,349 @@ +import { Module } from "./Module"; +import { ClassDef } from "./ClassDef"; +import { FunctionDef } from "./FunctionDef"; +import { Parameter } from "./Parameter"; +import { ClassReference } from "./ClassReference"; +import { Decorator } from "./Decorator"; +import { CodeBlock } from "./CodeBlock"; +import { Return } from "./Return"; +import { Assign } from "./Assign"; +import { Import } from "./Import"; + +describe("Python AST Integration", () => { + it("should generate a complete, realistic Python file", () => { + // Create the module with docstring + const module = new Module({ + name: "user_model", + docstring: + "Module containing the User data model and related functionality.", + }); + + // Add imports in a single block to avoid duplicates + const imports = [ + new Import({ moduleName: "typing", names: ["List", "Optional", "Dict"] }), + new Import({ moduleName: "datetime", names: ["datetime"] }), + new Import({ moduleName: "uuid", names: ["UUID"] }), + new Import({ moduleName: "dataclasses", names: ["dataclass", "field"] }), + ]; + + imports.forEach((imp) => module.addImport(imp)); + + // Create a data model class + const userClass = new ClassDef({ + name: "User", + decorators: [ + new Decorator({ name: "dataclass", moduleName: "dataclasses" }), + ], + docstring: + "Represents a user in the system with their associated metadata.", + }); + + // Add class variables with type annotations + userClass.addAttribute( + new Assign({ + target: "id", + type: new ClassReference({ name: "UUID", moduleName: "uuid" }), + value: "UUID.uuid4()", + isClassVariable: true, + }), + ); + + userClass.addAttribute( + new Assign({ + target: "username", + type: new ClassReference({ name: "str" }), + value: "''", + isClassVariable: true, + }), + ); + + userClass.addAttribute( + new Assign({ + target: "email", + type: new ClassReference({ name: "str" }), + value: "''", + isClassVariable: true, + }), + ); + + userClass.addAttribute( + new Assign({ + target: "created_at", + type: new ClassReference({ name: "datetime", moduleName: "datetime" }), + value: "datetime.now()", + isClassVariable: true, + }), + ); + + userClass.addAttribute( + new Assign({ + target: "roles", + type: new ClassReference({ name: "List", moduleName: "typing" }), + value: "field(default_factory=list)", + isClassVariable: true, + }), + ); + + userClass.addAttribute( + new Assign({ + target: "metadata", + type: new ClassReference({ + name: "Dict", + moduleName: "typing", + }), + value: "field(default_factory=dict)", + isClassVariable: true, + }), + ); + + // Add a class method for creating a user + const createMethod = new FunctionDef({ + name: "create", + isClassMethod: true, + parameters: [ + new Parameter({ name: "cls" }), + new Parameter({ + name: "username", + type: new ClassReference({ name: "str" }), + }), + new Parameter({ + name: "email", + type: new ClassReference({ name: "str" }), + }), + new Parameter({ + name: "roles", + type: new ClassReference({ + name: "List", + moduleName: "typing", + }), + default_: "None", + }), + ], + returnType: new ClassReference({ name: "User" }), + }); + + createMethod.addStatement( + new CodeBlock({ + code: `return cls( + id=UUID.uuid4(), + username=username, + email=email, + roles=roles or [], + created_at=datetime.now(), + metadata={})`, + }), + ); + + userClass.addMethod(createMethod); + + // Add an instance method + const addRoleMethod = new FunctionDef({ + name: "add_role", + parameters: [ + new Parameter({ name: "self" }), + new Parameter({ + name: "role", + type: new ClassReference({ name: "str" }), + }), + ], + returnType: new ClassReference({ name: "None" }), + }); + + addRoleMethod.addStatement( + new CodeBlock({ + code: "if role not in self.roles:\n self.roles.append(role)", + }), + ); + + userClass.addMethod(addRoleMethod); + + // Add a property + const isAdminProperty = new FunctionDef({ + name: "is_admin", + decorators: [new Decorator({ name: "property" })], + parameters: [new Parameter({ name: "self" })], + returnType: new ClassReference({ name: "bool" }), + }); + + isAdminProperty.addStatement( + new Return({ value: "'admin' in self.roles" }), + ); + + userClass.addMethod(isAdminProperty); + + // Add the class to the module + module.addClass(userClass); + + // Add example usage as a comment at the end of the file + module.addCodeBlock( + new CodeBlock({ + code: ` + +# Example usage: +if __name__ == "__main__": + user = User.create("john_doe", "john@example.com") + user.add_role("admin") + print(f"Is admin: {user.is_admin}") +`, + }), + ); + + // Generate the complete file + const output = module.toString(); + + // Compare with snapshot + expect(output).toMatchSnapshot(); + + // Additional specific assertions to ensure correct structure + const lines = output.split("\n"); + + // Check imports are at the top and not duplicated + const importLines = lines.filter( + (line) => line.startsWith("from") || line.startsWith("import"), + ); + const uniqueImports = new Set(importLines); + expect(importLines.length).toBe(uniqueImports.size); + + // Check docstring is after imports + const docstringIndex = lines.findIndex((line) => + line.includes('"""Module containing'), + ); + const lastImportIndex = lines.reduce( + (max, line, index) => + line.startsWith("from") || line.startsWith("import") ? index : max, + -1, + ); + expect(docstringIndex).toBeGreaterThan(lastImportIndex); + + // Check class is after docstring + const classIndex = lines.findIndex((line) => line.includes("class User:")); + expect(classIndex).toBeGreaterThan(docstringIndex); + + // Check example is at the end + const exampleIndex = lines.findIndex((line) => + line.includes("# Example usage:"), + ); + expect(exampleIndex).toBeGreaterThan(classIndex); + + // Check for specific content + expect(output).toContain("from typing import Dict, List, Optional"); + expect(output).toContain("from datetime import datetime"); + expect(output).toContain("from uuid import UUID"); + expect(output).toContain("from dataclasses import dataclass, field"); + expect(output).toContain("@dataclass"); + expect(output).toContain("class User:"); + expect(output).toContain("def create("); + expect(output).toContain("def add_role("); + expect(output).toContain("@property"); + expect(output).toContain("def is_admin("); + }); + + describe("Complex Class Generation", () => { + it("should generate a valid Python class with methods and type hints", () => { + // Create a new module + const module = new Module({ + name: "example", + docstring: "Example module demonstrating Python AST generation", + }); + + // Create class references for type hints + const strType = new ClassReference({ name: "str" }); + const intType = new ClassReference({ name: "int" }); + const optionalType = new ClassReference({ + name: "Optional", + moduleName: "typing", + }); + + // Create a base class + const baseClass = new ClassReference({ + name: "BaseClass", + moduleName: "base_module", + }); + + // Create the main class + const personClass = new ClassDef({ + name: "Person", + bases: [baseClass], + docstring: "A class representing a person with name and age.", + }); + + // Add class attributes + personClass.addAttribute( + new Assign({ + target: "default_age", + value: "0", + type: intType, + isClassVariable: true, + }), + ); + + // Add constructor + const initMethod = new FunctionDef({ + name: "__init__", + parameters: [ + new Parameter({ name: "self" }), + new Parameter({ name: "name", type: strType }), + new Parameter({ + name: "age", + type: optionalType, + default_: "None", + }), + ], + docstring: "Initialize a Person instance.", + }); + + // Add constructor body + initMethod.addStatement( + new CodeBlock({ + code: "self.name = name\nself.age = age if age is not None else self.default_age", + }), + ); + + personClass.addMethod(initMethod); + + // Add a static method + const fromDictMethod = new FunctionDef({ + name: "from_dict", + isStatic: true, + parameters: [ + new Parameter({ + name: "data", + type: new ClassReference({ name: "dict" }), + }), + ], + returnType: new ClassReference({ name: "Person" }), + docstring: "Create a Person instance from a dictionary.", + }); + + fromDictMethod.addStatement( + new CodeBlock({ + code: 'return Person(data["name"], data.get("age"))', + }), + ); + + personClass.addMethod(fromDictMethod); + + // Add an instance method with a decorator + const getInfoMethod = new FunctionDef({ + name: "get_info", + parameters: [new Parameter({ name: "self" })], + returnType: strType, + decorators: [new Decorator({ name: "property" })], + }); + + getInfoMethod.addStatement( + new CodeBlock({ + code: 'return f"{self.name} ({self.age} years old)"', + }), + ); + + personClass.addMethod(getInfoMethod); + + // Add the class to the module + module.addClass(personClass); + + // Generate the Python code and match against snapshot + expect(module.toString()).toMatchSnapshot(); + }); + }); +}); diff --git a/libs/python-ast/src/lib/core/AstNode.ts b/libs/python-ast/src/lib/core/AstNode.ts new file mode 100644 index 0000000..58bf221 --- /dev/null +++ b/libs/python-ast/src/lib/core/AstNode.ts @@ -0,0 +1,34 @@ +import { Writer } from "./Writer"; +import { IAstNode, IWriter } from "@amplication/ast-types"; + +/** + * Base class for all AST nodes in the Python AST library. + * This abstract class provides the foundation for all Python code generation nodes. + * Each specific node type must extend this class and implement the write method. + * + * @abstract + * @implements {IAstNode} + */ +export abstract class AstNode implements IAstNode { + /** + * Writes the AST node to the given writer. + * This method must be implemented by all concrete AST node classes. + * It defines how the node should be written to the output. + * + * @param {IWriter} writer - The writer instance to write the node's content + * @abstract + */ + public abstract write(writer: IWriter): void; + + /** + * Returns the string representation of the AST node. + * This method creates a new Writer instance and writes the node's content to it. + * + * @returns {string} The string representation of the AST node + */ + public toString(): string { + const writer = new Writer({}); + this.write(writer); + return writer.toString(); + } +} diff --git a/libs/python-ast/src/lib/core/Writer.ts b/libs/python-ast/src/lib/core/Writer.ts new file mode 100644 index 0000000..4da2f75 --- /dev/null +++ b/libs/python-ast/src/lib/core/Writer.ts @@ -0,0 +1,133 @@ +import { IWriter, IAstNode } from "@amplication/ast-types"; +import { Import } from "../ast/Import"; + +export interface WriterArgs { + currentModuleName?: string; +} + +/** + * Writer class for generating Python code from AST nodes. + */ +export class Writer implements IWriter { + private buffer: string[] = []; + private indentLevel = 0; + private readonly indentString = " "; + private lastCharacterIsNewline = true; + private imports: Set = new Set(); + private currentModuleName?: string; + + constructor(args: WriterArgs = {}) { + this.currentModuleName = args.currentModuleName; + } + + write(text: string): void { + if (!text) return; + + // Split text into lines and handle each line's indentation + const lines = text.split("\n"); + for (let i = 0; i < lines.length; i++) { + const line = lines[i]; + + // Add indentation at the start of each line if needed + if (this.lastCharacterIsNewline && line.length > 0) { + this.buffer.push(this.indentString.repeat(this.indentLevel)); + } + + // Write the line content + if (line.length > 0) { + this.buffer.push(line); + } + + // Add newline if this isn't the last line + if (i < lines.length - 1) { + this.buffer.push("\n"); + this.lastCharacterIsNewline = true; + } else { + this.lastCharacterIsNewline = text.endsWith("\n"); + } + } + } + + writeNode(node: IAstNode): void { + node.write(this); + } + + writeLine(text = ""): void { + this.write(text); + this.writeNewLineIfLastLineNot(); + } + + newLine(): void { + this.write("\n"); + } + + writeNewLineIfLastLineNot(): void { + if (!this.lastCharacterIsNewline) { + this.newLine(); + } + } + + indent(): void { + this.indentLevel++; + } + + dedent(): void { + if (this.indentLevel > 0) { + this.indentLevel--; + } + } + + addImport(node: Import): void { + this.imports.add(node); + } + + private stringifyImports(): string { + const importStatements: string[] = []; + const fromImports: Map> = new Map(); + + // Process all imports + for (const node of this.imports) { + const moduleName = node.moduleName; + if (!moduleName || moduleName === this.currentModuleName) continue; + + if (node.names.length === 0) { + // Simple import + importStatements.push( + `import ${moduleName}${node.alias ? ` as ${node.alias}` : ""}`, + ); + } else { + // From import + const moduleImports = fromImports.get(moduleName) || new Set(); + // Only apply alias if there's exactly one name being imported + if (node.names.length === 1) { + node.names.forEach((name) => { + moduleImports.add(name + (node.alias ? ` as ${node.alias}` : "")); + }); + } else { + node.names.forEach((name) => { + moduleImports.add(name); + }); + } + fromImports.set(moduleName, moduleImports); + } + } + + // Sort and add from imports + Array.from(fromImports.entries()) + .sort(([a], [b]) => a.localeCompare(b)) + .forEach(([moduleName, names]) => { + const sortedNames = Array.from(names).sort(); + importStatements.push( + `from ${moduleName} import ${sortedNames.join(", ")}`, + ); + }); + + return importStatements.length > 0 + ? importStatements.join("\n") + "\n\n" + : ""; + } + + toString(): string { + return this.stringifyImports() + this.buffer.join(""); + } +} diff --git a/libs/python-ast/src/lib/core/index.ts b/libs/python-ast/src/lib/core/index.ts new file mode 100644 index 0000000..1442c4c --- /dev/null +++ b/libs/python-ast/src/lib/core/index.ts @@ -0,0 +1,2 @@ +export * from "./AstNode"; +export * from "./Writer"; diff --git a/libs/python-ast/tsconfig.json b/libs/python-ast/tsconfig.json new file mode 100644 index 0000000..f5b8565 --- /dev/null +++ b/libs/python-ast/tsconfig.json @@ -0,0 +1,22 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "module": "commonjs", + "forceConsistentCasingInFileNames": true, + "strict": true, + "noImplicitOverride": true, + "noPropertyAccessFromIndexSignature": true, + "noImplicitReturns": true, + "noFallthroughCasesInSwitch": true + }, + "files": [], + "include": [], + "references": [ + { + "path": "./tsconfig.lib.json" + }, + { + "path": "./tsconfig.spec.json" + } + ] +} diff --git a/libs/python-ast/tsconfig.lib.json b/libs/python-ast/tsconfig.lib.json new file mode 100644 index 0000000..33eca2c --- /dev/null +++ b/libs/python-ast/tsconfig.lib.json @@ -0,0 +1,10 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "outDir": "../../dist/out-tsc", + "declaration": true, + "types": ["node"] + }, + "include": ["src/**/*.ts"], + "exclude": ["jest.config.ts", "src/**/*.spec.ts", "src/**/*.test.ts"] +} diff --git a/libs/python-ast/tsconfig.spec.json b/libs/python-ast/tsconfig.spec.json new file mode 100644 index 0000000..9b2a121 --- /dev/null +++ b/libs/python-ast/tsconfig.spec.json @@ -0,0 +1,14 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "outDir": "../../dist/out-tsc", + "module": "commonjs", + "types": ["jest", "node"] + }, + "include": [ + "jest.config.ts", + "src/**/*.test.ts", + "src/**/*.spec.ts", + "src/**/*.d.ts" + ] +} diff --git a/typedoc.json b/typedoc.json index b04bf50..b9d6cb0 100644 --- a/typedoc.json +++ b/typedoc.json @@ -1,6 +1,11 @@ { "entryPointStrategy": "packages", - "entryPoints": ["libs/ast-types", "libs/csharp-ast", "libs/java-ast"], + "entryPoints": [ + "libs/ast-types", + "libs/csharp-ast", + "libs/java-ast", + "libs/python-ast" + ], "theme": "default", "name": "Amplication AST Libraries Documentation", "tsconfig": "tsconfig.base.json",