Python context (extend + test)

core/autocomplete/context/root-path-context/RootPathContextService.ts

@@ -26,8 +26,10 @@ export class RootPathContextService {
   private static TYPES_TO_USE = new Set([
     "program",
     "function_declaration",
+    "function_definition",
     "method_definition",
     "class_declaration",
+    "class_definition",
   ]);
 
   /**
@@ -97,6 +99,7 @@ export class RootPathContextService {
         contents: await this.ide.readRangeInFile(def.filepath, def.range),
       })),
     );
+
     return newSnippets;
   }
 

core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts

@@ -2,41 +2,62 @@ import { testRootPathContext } from "./testUtils";
 
 const TEST_CASES = [
   {
-    description: "function",
+    nodeType: "function_declaration",
     fileName: "file1.ts",
-    range: {
-      start: { line: 10, character: 2 },
-      end: { line: 10, character: 24 },
-    },
-    positions: [
+    language: "TypeScript",
+    cursorPosition: { line: 10, character: 24 },
+    definitionPositions: [
       { row: 9, column: 34 }, // Person
       { row: 9, column: 44 }, // Address
     ],
   },
   {
-    description: "class method",
+    nodeType: "method_declaration",
     fileName: "file1.ts",
-    range: {
-      start: { line: 22, character: 4 },
-      end: { line: 22, character: 30 },
-    },
-    positions: [
+    language: "TypeScript",
+    cursorPosition: { line: 22, character: 30 },
+    definitionPositions: [
       { row: 13, column: 29 }, // BaseClass
       { row: 13, column: 55 }, // FirstInterface
       { row: 13, column: 72 }, // SecondInterface
       { row: 21, column: 33 }, // Person
       { row: 21, column: 43 }, // Address
     ],
   },
+  {
+    nodeType: "function_definition",
+    fileName: "file1.py",
+    language: "Python",
+    cursorPosition: { line: 4, character: 25 },
+    definitionPositions: [
+      { row: 3, column: 30 }, // Person
+      { row: 3, column: 42 }, // Address
+    ],
+  },
+  {
+    nodeType: "function_definition (inside a class)",
+    fileName: "file1.py",
+    language: "Python",
+    cursorPosition: { line: 12, character: 33 },
+    definitionPositions: [
+      { row: 6, column: 21 }, // BaseClass
+      { row: 6, column: 33 }, // Collection
+      { row: 11, column: 47 }, // Person
+      { row: 11, column: 59 }, // Address
+    ],
+  },
 ];
 
 describe("RootPathContextService", () => {
-  describe("TypeScript should return expected snippets when editing inside a:", () => {
-    test.each(TEST_CASES)(
-      "should look for correct type definitions when editing inside a $description",
-      async ({ fileName, range, positions }) => {
-        await testRootPathContext("typescript", fileName, range, positions);
-      },
-    );
-  });
+  test.each(TEST_CASES)(
+    "Should look for correct type definitions when editing inside a $nodeType in $language",
+    async ({ fileName, cursorPosition, definitionPositions }) => {
+      await testRootPathContext(
+        "files",
+        fileName,
+        cursorPosition,
+        definitionPositions,
+      );
+    },
+  );
 });

core/autocomplete/context/root-path-context/test/files/base_module.py

@@ -0,0 +1,26 @@
+# File: base_module.py
+
+class BaseClass:
+    def __init__(self):
+        print("BaseClass initialized")
+
+class Collection:
+    def __init__(self):
+        print("Collection initialized")
+
+class Address:
+    def __init__(self, street: str, city: str, zip_code: str):
+        self.street = street
+        self.city = city
+        self.zip_code = zip_code
+
+    def __str__(self):
+        return f"{self.street}, {self.city}, {self.zip_code}"
+
+class Person:
+    def __init__(self, name: str, address: Address):
+        self.name = name
+        self.address = address
+
+    def __str__(self):
+        return f"{self.name} lives at {self.address}"

core/autocomplete/context/root-path-context/test/files/file1.py

@@ -0,0 +1,13 @@
+from .base_module import BaseClass, Collection, Person, Address
+from typing import List
+
+def get_address(person: Person) -> Address:
+    return person.address
+
+class Group(BaseClass, Collection):
+    def __init__(self, people: List[Person]) -> None:
+        super().__init__()
+        self.people = people
+
+    def get_person_address(self, person: Person) -> Address:
+        return get_address(person)

core/autocomplete/context/root-path-context/test/testUtils.ts

@@ -1,48 +1,34 @@
 import { jest } from "@jest/globals";
 import fs from "fs";
 import path from "path";
 
 import Parser from "web-tree-sitter";
-import { Range } from "../../../..";
+import { Position } from "../../../..";
 import { testIde } from "../../../../test/util/fixtures";
 import { getAst, getTreePathAtCursor } from "../../../util/ast";
 import { ImportDefinitionsService } from "../../ImportDefinitionsService";
 import { RootPathContextService } from "../RootPathContextService";
 
-function splitTextAtRange(fileContent: string, range: Range): [string, string] {
+function splitTextAtPosition(
+  fileContent: string,
+  position: Position,
+): [string, string] {
   const lines = fileContent.split("\n");
   let currentPos = 0;
 
-  if (range.start.line === range.end.line) {
-    // If range is on the same line, calculate position once
-    for (let i = 0; i < range.start.line; i++) {
-      currentPos += lines[i].length + 1; // +1 for the newline character
-    }
-    const startPos = currentPos + range.start.character;
-    const endPos = currentPos + range.end.character;
-    return [fileContent.slice(0, startPos), fileContent.slice(endPos)];
+  // Calculate position based on the provided line and character
+  for (let i = 0; i < position.line; i++) {
+    currentPos += lines[i].length + 1; // +1 for the newline character
   }
+  const splitPos = currentPos + position.character;
 
-  // Calculate position of range start
-  for (let i = 0; i < range.start.line; i++) {
-    currentPos += lines[i].length + 1;
-  }
-  const startPos = currentPos + range.start.character;
-
-  // Calculate position of range end
-  currentPos = 0;
-  for (let i = 0; i < range.end.line; i++) {
-    currentPos += lines[i].length + 1;
-  }
-  const endPos = currentPos + range.end.character;
-
-  return [fileContent.slice(0, startPos), fileContent.slice(endPos)];
+  return [fileContent.slice(0, splitPos), fileContent.slice(splitPos)];
 }
 
 export async function testRootPathContext(
   folderName: string,
   relativeFilepath: string,
-  rangeToFill: Range,
+  position: Position,
   expectedDefinitionPositions: Parser.Point[],
 ) {
   // Create a mocked instance of RootPathContextService
@@ -76,9 +62,9 @@
 
   // Get results of root path context
   const startPath = path.join(testFolderPath, relativeFilepath);
-  const [prefix, suffix] = splitTextAtRange(
+  const [prefix, suffix] = splitTextAtPosition(
     fs.readFileSync(startPath, "utf8"),
-    rangeToFill,
+    position,
   );
   const fileContents = prefix + suffix;
   const ast = await getAst(startPath, fileContents);

extensions/vscode/tree-sitter/root-path-context-queries/class_definition/python.scm

@@ -0,0 +1,5 @@
+(
+  (class_definition
+    superclasses: (argument_list
+      (identifier) @superclass))
+)