diff --git a/core/autocomplete/context/root-path-context/RootPathContextService.ts b/core/autocomplete/context/root-path-context/RootPathContextService.ts index b381f2166..accd4f6df 100644 --- a/core/autocomplete/context/root-path-context/RootPathContextService.ts +++ b/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; } diff --git a/core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts b/core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts index 597d37d12..196a7b643 100644 --- a/core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts +++ b/core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts @@ -2,25 +2,21 @@ 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 @@ -28,15 +24,40 @@ const TEST_CASES = [ { 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, + ); + }, + ); }); diff --git a/core/autocomplete/context/root-path-context/test/files/__init__.py b/core/autocomplete/context/root-path-context/test/files/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/core/autocomplete/context/root-path-context/test/files/base_module.py b/core/autocomplete/context/root-path-context/test/files/base_module.py new file mode 100644 index 000000000..cf0dc0abb --- /dev/null +++ b/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}" diff --git a/core/autocomplete/context/root-path-context/test/files/file1.py b/core/autocomplete/context/root-path-context/test/files/file1.py new file mode 100644 index 000000000..55c3d9c25 --- /dev/null +++ b/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) diff --git a/core/autocomplete/context/root-path-context/test/typescript/file1.ts b/core/autocomplete/context/root-path-context/test/files/file1.ts similarity index 100% rename from core/autocomplete/context/root-path-context/test/typescript/file1.ts rename to core/autocomplete/context/root-path-context/test/files/file1.ts diff --git a/core/autocomplete/context/root-path-context/test/testUtils.ts b/core/autocomplete/context/root-path-context/test/testUtils.ts index 503383fd2..9526cacdb 100644 --- a/core/autocomplete/context/root-path-context/test/testUtils.ts +++ b/core/autocomplete/context/root-path-context/test/testUtils.ts @@ -3,46 +3,32 @@ 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 @@ export async function testRootPathContext( // 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); diff --git a/extensions/vscode/tree-sitter/root-path-context-queries/class_definition/python.scm b/extensions/vscode/tree-sitter/root-path-context-queries/class_definition/python.scm new file mode 100644 index 000000000..5a9e34b86 --- /dev/null +++ b/extensions/vscode/tree-sitter/root-path-context-queries/class_definition/python.scm @@ -0,0 +1,5 @@ +( + (class_definition + superclasses: (argument_list + (identifier) @superclass)) +) \ No newline at end of file