Parser Role

src/app.ts

@@ -1,3 +1,4 @@
+
 import { Octokit } from "@octokit/rest";
 import { createNodeMiddleware } from "@octokit/webhooks";
 import { WebhookEventMap } from "@octokit/webhooks-definitions/schema";
@@ -46,8 +47,6 @@ async function handlePullRequestOpened({
   console.log(
     `Received a pull request event for #${payload.pull_request.number}`
   );
-  // const reposWithInlineEnabled = new Set<number>([601904706, 701925328]);
-  // const canInlineSuggest = reposWithInlineEnabled.has(payload.repository.id);
   try {
     console.log("pr info", {
       id: payload.repository.id,
@@ -69,7 +68,6 @@ async function handlePullRequestOpened({
 }
 
 // This sets up a webhook event listener. When your app receives a webhook event from GitHub with a `X-GitHub-Event` header value of `pull_request` and an `action` payload value of `opened`, it calls the `handlePullRequestOpened` event handler that is defined above.
-//@ts-ignore
 reviewApp.webhooks.on("pull_request.opened", handlePullRequestOpened);
 
 const port = process.env.PORT || 3000;

src/constants.ts

@@ -1,6 +1,8 @@
 import { Node } from "@babel/traverse";
 import { JavascriptParser } from "./context/language/javascript-parser";
 import { ChatCompletionMessageParam } from "groq-sdk/resources/chat/completions";
+import Parser from "tree-sitter";
+import { PythonParser } from "./context/language/python-parser";
 
 export interface PRFile {
   sha: string;
@@ -89,8 +91,16 @@ export const processGitFilepath = (filepath: string) => {
   return filepath.startsWith("/") ? filepath.slice(1) : filepath;
 };
 
+export const isBabelNode = (node: Node | Parser.SyntaxNode): node is Node  =>{
+  if (node === null) {
+    return false;
+  }
+
+  return (node as Parser.SyntaxNode).childCount === undefined;
+}
+
 export interface EnclosingContext {
-  enclosingContext: Node | null;
+  enclosingContext: Node | Parser.SyntaxNode | null;
 }
 
 export interface AbstractParser {
@@ -107,6 +117,7 @@ const EXTENSIONS_TO_PARSERS: Map<string, AbstractParser> = new Map([
   ["tsx", new JavascriptParser()],
   ["js", new JavascriptParser()],
   ["jsx", new JavascriptParser()],
+  ["py", new PythonParser]
 ]);
 
 export const getParserForExtension = (filename: string) => {
@@ -123,4 +134,4 @@ export const assignLineNumbers = (contents: string): string => {
     return numberedLine;
   });
   return linesWithNumbers.join("\n");
-};
+};

src/context/language/javascript-parser.ts

@@ -74,4 +74,4 @@ export class JavascriptParser implements AbstractParser {
       };
     }
   }
-}
+}

src/context/language/python-parser.ts

@@ -1,15 +1,69 @@
 import { AbstractParser, EnclosingContext } from "../../constants";
+import Parser from "tree-sitter";
+import Python from "tree-sitter-python";
+
+const parser = new Parser();
+parser.setLanguage(Python);
+
+const processNode = (
+  node: Parser.SyntaxNode,
+  lineStart: number,
+  lineEnd: number,
+  largestSize: number,
+  largestEnclosingContext: Parser.SyntaxNode | null
+) => {
+  const { startPosition, endPosition } = node;
+  if (startPosition.row <= lineStart && lineEnd <= endPosition.row) {
+    const size = endPosition.row - startPosition.row;
+    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;
+    let largestSize = 0;
+    let largestEnclosingContext: Parser.SyntaxNode = null;
+
+    const ast = parser.parse(file);
+    const rootNode = ast.rootNode;
+
+    const traverse = (node: Parser.SyntaxNode) => {
+
+      if (node.type == "function_definition") {
+        const result = processNode(node, lineStart, lineEnd, largestSize, largestEnclosingContext)
+        largestSize = result.largestSize;
+        largestEnclosingContext = result.largestEnclosingContext;
+      }
+
+      for (let i = 0; i < node.childCount; i++)
+      {
+        const child = node.child(i);
+        if (child?.childCount > 0) {
+          traverse(child)
+        }
+      }
+    }
+    traverse(rootNode)
+
+    return {
+      enclosingContext: largestEnclosingContext
+    } as EnclosingContext;
   }
   dryRun(file: string): { valid: boolean; error: string } {
-    // TODO: Implement this method for Python
-    return { valid: false, error: "Not implemented yet" };
+    try {
+      const ast = parser.parse(file)
+      return { valid: true, error: ""}
+    }
+    catch (exc) {
+      return { valid: false, error: exc };
+    }
   }
-}
+}

src/context/review.ts

@@ -1,11 +1,15 @@
 import {
   AbstractParser,
+  EnclosingContext,
   PRFile,
   PatchInfo,
   getParserForExtension,
+  isBabelNode,
 } from "../constants";
 import * as diff from "diff";
 import { JavascriptParser } from "./language/javascript-parser";
+import { PythonParser } from "./language/python-parser";
+import Parser from "tree-sitter";
 import { Node } from "@babel/traverse";
 
 const expandHunk = (
@@ -107,13 +111,20 @@ const trimHunk = (hunk: diff.Hunk): diff.Hunk => {
 
 const buildingScopeString = (
   currentFile: string,
-  scope: Node,
+  scope: EnclosingContext,
   hunk: diff.Hunk
 ) => {
   const res: string[] = [];
   const trimmedHunk = trimHunk(hunk);
-  const functionStartLine = scope.loc.start.line;
-  const functionEndLine = scope.loc.end.line;
+  let functionStartLine = 0;
+  let functionEndLine = 0;
+  if (isBabelNode(scope.enclosingContext)) {
+    functionStartLine = scope.enclosingContext.loc.start.line;
+    functionEndLine = scope.enclosingContext.loc.end.line;
+  } else {
+    functionStartLine = scope.enclosingContext.startPosition.row;
+    functionEndLine = scope.enclosingContext.endPosition.row;
+  }
   const updatedFileLines = currentFile.split("\n");
   // Extract the lines of the function
   const functionContext = updatedFileLines.slice(
@@ -196,7 +207,7 @@ const diffContextPerHunk = (file: PRFile, parser: AbstractParser) => {
   const hunks: diff.Hunk[] = [];
   const order: number[] = [];
   const scopeRangeHunkMap = new Map<string, diff.Hunk[]>();
-  const scopeRangeNodeMap = new Map<string, Node>();
+  const scopeRangeNodeMap = new Map<string, EnclosingContext>();
   const expandStrategy: diff.Hunk[] = [];
 
   patches.forEach((p) => {
@@ -205,6 +216,7 @@ const diffContextPerHunk = (file: PRFile, parser: AbstractParser) => {
     });
   });
 
+  console.info("processing hunks");
   hunks.forEach((hunk, idx) => {
     try {
       const trimmedHunk = trimHunk(hunk);
@@ -213,14 +225,30 @@ const diffContextPerHunk = (file: PRFile, parser: AbstractParser) => {
       ).length;
       const lineStart = trimmedHunk.newStart;
       const lineEnd = lineStart + insertions;
+      console.info(`searching for context for range: ${lineStart}-${lineEnd}`);
       const largestEnclosingFunction = parser.findEnclosingContext(
         updatedFile,
         lineStart,
         lineEnd
-      ).enclosingContext;
+      );
+      const node = largestEnclosingFunction.enclosingContext;
+
+      if (isBabelNode(node)) {
+        console.log("Babel Node");
+      }
+      else {
+        console.log("Python Node")
+      }
 
       if (largestEnclosingFunction) {
-        const enclosingRangeKey = `${largestEnclosingFunction.loc.start.line} -> ${largestEnclosingFunction.loc.end.line}`;
+        let enclosingRangeKey = "";
+        if (isBabelNode(node)) {
+          enclosingRangeKey = `${node.loc.start.line} -> ${node.loc.end.line}`;
+          console.log('enclosed context rarnge for Babel Node: ', enclosingRangeKey)
+        } else {
+          enclosingRangeKey = `${node.startPosition.row} -> ${node.endPosition.row}`;
+          console.log('enclosed context rarnge for Python Node: ', enclosingRangeKey)
+        }
         let existingHunks = scopeRangeHunkMap.get(enclosingRangeKey) || [];
         existingHunks.push(hunk);
         scopeRangeHunkMap.set(enclosingRangeKey, existingHunks);
@@ -246,12 +274,23 @@ const diffContextPerHunk = (file: PRFile, parser: AbstractParser) => {
 
   const contexts: string[] = [];
   scopeStategy.forEach(([rangeKey, hunk]) => {
-    const context = buildingScopeString(
-      updatedFile,
-      scopeRangeNodeMap.get(rangeKey),
-      hunk
-    ).join("\n");
-    contexts.push(context);
+    const enclosingContext = scopeRangeNodeMap.get(rangeKey);
+    const node = enclosingContext.enclosingContext;
+    if (node && isBabelNode(node)) {
+      const context = buildingScopeString(
+        updatedFile,
+        enclosingContext,
+        hunk
+      ).join("\n");
+      contexts.push(context);
+    } else if (node) {
+      const context = buildingScopeString(
+        updatedFile,
+        enclosingContext,
+        hunk
+      ).join("\n");
+      contexts.push(context);
+    }
   });
   expandStrategy.forEach((hunk) => {
     const context = expandHunk(file.old_contents, hunk);
@@ -275,11 +314,19 @@ const functionContextPatchStrategy = (
   return res;
 };
 
+
 export const smarterContextPatchStrategy = (file: PRFile) => {
   const parser: AbstractParser = getParserForExtension(file.filename);
   if (parser != null) {
+    if (parser instanceof PythonParser) {
+      console.log('Using PythonParser')
+    }
+    else if (parser instanceof JavascriptParser) {
+      console.log('Using JavaScriptParser')
+    }
     return functionContextPatchStrategy(file, parser);
   } else {
+    console.info("Using basic parser");
     return expandedPatchStrategy(file);
   }
-};
+};

src/env.ts

@@ -4,13 +4,15 @@ import chalk from "chalk";
 
 dotenv.config();
 
+
 export const env = {
   GITHUB_APP_ID: process.env.GITHUB_APP_ID,
   GITHUB_PRIVATE_KEY: process.env.GITHUB_PRIVATE_KEY,
   GITHUB_WEBHOOK_SECRET: process.env.GITHUB_WEBHOOK_SECRET,
   GROQ_API_KEY: process.env.GROQ_API_KEY,
 } as const;
 
+
 let valid = true;
 
 for (const key in env) {

src/prompts.ts

@@ -166,4 +166,4 @@ export const isConversationWithinLimit = (
   // the one for gpt-3.5-turbo as a rough equivalent.
   const convoTokens = encodeChat(convo, "gpt-3.5-turbo").length;
   return convoTokens < ModelsToTokenLimits[model];
-};
+};