Add Python parser to support .py files

jest.config.js

@@ -0,0 +1,5 @@
+module.exports = {
+  preset: "ts-jest", // Enables TypeScript support for Jest
+  testEnvironment: "node", // Simulates a Node.js environment for tests
+  testMatch: ["<rootDir>/tests/**/*.test.ts", "<rootDir>/tests/**/*.spec.ts"], // Matches test files in the "tests" directory
+};

package.json

@@ -5,6 +5,7 @@
   "license": "ISC",
   "scripts": {
     "start": "tsx src/app.ts"
+    "test": "jest"
   },
   "dependencies": {
     "@babel/parser": "^7.23.0",

src/constants.ts

@@ -1,5 +1,6 @@
 import { Node } from "@babel/traverse";
 import { JavascriptParser } from "./context/language/javascript-parser";
+import { PythonParser } from "./context/language/python-parser"; // Import PythonParser
 import { ChatCompletionMessageParam } from "groq-sdk/resources/chat/completions";
 
 export interface PRFile {
@@ -107,6 +108,7 @@ const EXTENSIONS_TO_PARSERS: Map<string, AbstractParser> = new Map([
   ["tsx", new JavascriptParser()],
   ["js", new JavascriptParser()],
   ["jsx", new JavascriptParser()],
+  ["py", new PythonParser()], // Register PythonParser
 ]);
 
 export const getParserForExtension = (filename: string) => {
@@ -124,3 +126,4 @@ export const assignLineNumbers = (contents: string): string => {
   });
   return linesWithNumbers.join("\n");
 };
+

src/context/language/python-parser.ts

@@ -1,15 +1,54 @@
 import { AbstractParser, EnclosingContext } from "../../constants";
+import Parser from "tree-sitter";
+import Python from "tree-sitter-python";
+
 export class PythonParser implements AbstractParser {
+  parser: Parser;
+
+  constructor() {
+    this.parser = new Parser();
+    this.parser.setLanguage(Python);
+  }
+
+  dryRun(file: string): { valid: boolean; error: string } {
+    try {
+      this.parser.parse(file);
+      return { valid: true, error: "" };
+    } catch (error) {
+      return { valid: false, error: error.message };
+    }
+  }
+
   findEnclosingContext(
     file: string,
     lineStart: number,
     lineEnd: number
   ): EnclosingContext {
-    // TODO: Implement this method for Python
-    return null;
-  }
-  dryRun(file: string): { valid: boolean; error: string } {
-    // TODO: Implement this method for Python
-    return { valid: false, error: "Not implemented yet" };
+    const tree = this.parser.parse(file);
+    let largestNode: Parser.SyntaxNode = null;
+    let largestSize = 0;
+
+    const visitNode = (node: Parser.SyntaxNode) => {
+      const nodeStart = node.startPosition.row + 1;
+      const nodeEnd = node.endPosition.row + 1;
+
+      if (nodeStart <= lineStart && lineEnd <= nodeEnd) {
+        const size = nodeEnd - nodeStart;
+        if (size > largestSize) {
+          largestSize = size;
+          largestNode = node;
+        }
+      }
+
+      for (const child of node.namedChildren) {
+        visitNode(child);
+      }
+    };
+
+    visitNode(tree.rootNode);
+
+    return {
+      enclosingContext: largestNode ? largestNode.type : null,
+    };
   }
 }

tests/python-parser.test.ts

@@ -0,0 +1,64 @@
+import { PythonParser } from "../src/context/language/python-parser";
+
+describe("PythonParser", () => {
+  const parser = new PythonParser();
+
+  test("dryRun returns valid for correct Python code", () => {
+    const validCode = `
+def foo():
+    pass
+`;
+    const result = parser.dryRun(validCode);
+    expect(result.valid).toBe(true);
+    expect(result.error).toBe("");
+  });
+
+  test("dryRun returns invalid for incorrect Python code", () => {
+    const invalidCode = `
+def foo(
+    pass
+`;
+    const result = parser.dryRun(invalidCode);
+    expect(result.valid).toBe(false);
+    expect(result.error).not.toBe("");
+  });
+
+  test("findEnclosingContext returns correct context for a function", () => {
+    const pythonCode = `
+class MyClass:
+    def method(self):
+        pass
+
+def foo():
+    pass
+`;
+    const context = parser.findEnclosingContext(pythonCode, 3, 3); // Line 3 corresponds to 'method'
+    expect(context.enclosingContext).toEqual("function_definition");
+  });
+
+  test("findEnclosingContext returns correct context for a class", () => {
+    const pythonCode = `
+class MyClass:
+    def method(self):
+        pass
+
+def foo():
+    pass
+`;
+    const context = parser.findEnclosingContext(pythonCode, 2, 3); // Line range includes the class
+    expect(context.enclosingContext).toEqual("class_definition");
+  });
+
+  test("findEnclosingContext returns null for unrelated lines", () => {
+    const pythonCode = `
+class MyClass:
+    def method(self):
+        pass
+
+def foo():
+    pass
+`;
+    const context = parser.findEnclosingContext(pythonCode, 10, 11); // Outside any context
+    expect(context.enclosingContext).toBeNull();
+  });
+});