diff --git a/package.json b/package.json index fb08515..e1fc99e 100644 --- a/package.json +++ b/package.json @@ -18,7 +18,8 @@ "gpt-tokenizer": "^2.1.2", "groq-sdk": "^0.8.0", "octokit": "^3.1.1", - "xml2js": "^0.6.2" + "xml2js": "^0.6.2", + "python-ast": "^0.1.0" }, "devDependencies": { "@types/babel__traverse": "^7.20.3", diff --git a/src/constants.ts b/src/constants.ts index 14c7de1..8b9ad78 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -1,7 +1,8 @@ import { Node } from "@babel/traverse"; import { JavascriptParser } from "./context/language/javascript-parser"; import { ChatCompletionMessageParam } from "groq-sdk/resources/chat/completions"; - +import { PythonParser } from "./context/language/python-parser"; +import { PythonNode } from "./context/language/python-parser"; export interface PRFile { sha: string; filename: string; @@ -90,7 +91,7 @@ export const processGitFilepath = (filepath: string) => { }; export interface EnclosingContext { - enclosingContext: Node | null; + enclosingContext: Node | PythonNode | null; } export interface AbstractParser { @@ -107,6 +108,8 @@ const EXTENSIONS_TO_PARSERS: Map = new Map([ ["tsx", new JavascriptParser()], ["js", new JavascriptParser()], ["jsx", new JavascriptParser()], + ["py", new PythonParser()], + ["pyx", new PythonParser()], ]); export const getParserForExtension = (filename: string) => { diff --git a/src/context/language/python-parser.ts b/src/context/language/python-parser.ts index 845e90b..ef9d2f6 100644 --- a/src/context/language/python-parser.ts +++ b/src/context/language/python-parser.ts @@ -1,15 +1,84 @@ import { AbstractParser, EnclosingContext } from "../../constants"; +import * as pythonAst from "python-ast"; + +export interface PythonNode { + type: string; + start: number; + end: number; + lineno: number; + end_lineno: number; + body?: PythonNode[]; +} + +const processNode = ( + node: PythonNode, + lineStart: number, + lineEnd: number, + largestSize: number, + largestEnclosingContext: PythonNode | null +) => { + if (node.lineno <= lineStart && lineEnd <= node.end_lineno) { + const size = node.end_lineno - node.lineno; + if (size > largestSize) { + largestSize = size; + largestEnclosingContext = node; + } + } + return { largestSize, largestEnclosingContext }; +}; + export class PythonParser implements AbstractParser { findEnclosingContext( file: string, lineStart: number, lineEnd: number ): EnclosingContext { - // TODO: Implement this method for Python - return null; + try { + const ast = pythonAst.parse(file); + let largestEnclosingContext: PythonNode = null; + let largestSize = 0; + + const visitNode = (node: PythonNode) => { + // Look for function definitions and class definitions + if (node.type === "FunctionDef" || node.type === "ClassDef") { + ({ largestSize, largestEnclosingContext } = processNode( + node, + lineStart, + lineEnd, + largestSize, + largestEnclosingContext + )); + } + + // Recursively visit child nodes + if (node.body) { + node.body.forEach(visitNode); + } + }; + + visitNode(ast as unknown as PythonNode); + + return { + enclosingContext: largestEnclosingContext, + } as EnclosingContext; + } catch (error) { + console.error("Error parsing Python file:", error); + return { enclosingContext: null }; + } } + dryRun(file: string): { valid: boolean; error: string } { - // TODO: Implement this method for Python - return { valid: false, error: "Not implemented yet" }; + try { + pythonAst.parse(file); + return { + valid: true, + error: "", + }; + } catch (exc) { + return { + valid: false, + error: exc.toString(), + }; + } } }