diff --git a/jest.config.js b/jest.config.js new file mode 100644 index 0000000..aa690d4 --- /dev/null +++ b/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: ["/tests/**/*.test.ts", "/tests/**/*.spec.ts"], // Matches test files in the "tests" directory +}; diff --git a/package.json b/package.json index fb08515..0b51c8c 100644 --- a/package.json +++ b/package.json @@ -5,6 +5,7 @@ "license": "ISC", "scripts": { "start": "tsx src/app.ts" + "test": "jest" }, "dependencies": { "@babel/parser": "^7.23.0", diff --git a/src/constants.ts b/src/constants.ts index 14c7de1..468dab7 100644 --- a/src/constants.ts +++ b/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 = 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"); }; + diff --git a/src/context/language/python-parser.ts b/src/context/language/python-parser.ts index 845e90b..266d7ea 100644 --- a/src/context/language/python-parser.ts +++ b/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, + }; } } diff --git a/tests/python-parser.test.ts b/tests/python-parser.test.ts new file mode 100644 index 0000000..f26cfc8 --- /dev/null +++ b/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(); + }); +});