Add Python Parser Implementation

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",

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<string, AbstractParser> = new Map([
   ["tsx", new JavascriptParser()],
   ["js", new JavascriptParser()],
   ["jsx", new JavascriptParser()],
+  ["py", new PythonParser()],
+  ["pyx", new PythonParser()],
 ]);
 
 export const getParserForExtension = (filename: string) => {

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(),
+      };
+    }
   }
 }