diff --git a/src/services/tree-sitter/__tests__/languageParser.spec.ts b/src/services/tree-sitter/__tests__/languageParser.spec.ts index 44811d46c6..9bf9ecc881 100644 --- a/src/services/tree-sitter/__tests__/languageParser.spec.ts +++ b/src/services/tree-sitter/__tests__/languageParser.spec.ts @@ -1,136 +1,56 @@ // npx vitest services/tree-sitter/__tests__/languageParser.spec.ts +import * as path from "path" import { loadRequiredLanguageParsers } from "../languageParser" -vi.mock("web-tree-sitter", () => { - const mockParserInit = vi.fn().mockResolvedValue(undefined) - const mockLanguageLoad = vi.fn().mockResolvedValue({ - query: vi.fn().mockReturnValue({ id: "mock-query" }), - }) - const mockSetLanguage = vi.fn() +// Path to the directory containing the WASM files. +const WASM_DIR = path.join(__dirname, "../../../node_modules/tree-sitter-wasms/out") - // Create a constructor function that also has static methods - function MockParser() { - return { - setLanguage: mockSetLanguage, - } - } - MockParser.init = mockParserInit - - return { - Parser: MockParser, - Language: { - load: mockLanguageLoad, - }, - // Export the mocks so tests can access them - __mocks: { - mockParserInit, - mockLanguageLoad, - mockSetLanguage, - }, - } -}) - -// Import the mocked module to get access to the mock functions -const { __mocks } = (await import("web-tree-sitter")) as any -const { mockParserInit, mockLanguageLoad, mockSetLanguage } = __mocks - -describe("Language Parser", () => { - beforeEach(() => { - vi.clearAllMocks() +describe("loadRequiredLanguageParsers", () => { + it("should load Python parser for .py files", async () => { + const files = ["test.py"] + const parsers = await loadRequiredLanguageParsers(files, WASM_DIR) + expect(parsers.py).toBeDefined() }) - describe("loadRequiredLanguageParsers", () => { - it("should initialize parser only once", async () => { - const files = ["test.js", "test2.js"] - await loadRequiredLanguageParsers(files) - await loadRequiredLanguageParsers(files) + it("should load JavaScript parser for .js and .jsx files", async () => { + const files = ["test.js", "test.jsx"] + const parsers = await loadRequiredLanguageParsers(files, WASM_DIR) + expect(parsers.js).toBeDefined() + expect(parsers.jsx).toBeDefined() + expect(parsers.js.query).toBeDefined() + expect(parsers.jsx.query).toBeDefined() + }) - expect(mockParserInit).toHaveBeenCalledTimes(1) - }) + it("should load multiple language parsers as needed", async () => { + const files = ["test.js", "test.py", "test.rs", "test.go"] + const parsers = await loadRequiredLanguageParsers(files, WASM_DIR) + expect(parsers.js).toBeDefined() + expect(parsers.py).toBeDefined() + expect(parsers.rs).toBeDefined() + expect(parsers.go).toBeDefined() + }) - it("should load JavaScript parser for .js and .jsx files", async () => { - const files = ["test.js", "test.jsx"] - const parsers = await loadRequiredLanguageParsers(files) + it("should handle C/C++ files correctly", async () => { + const files = ["test.c", "test.h", "test.cpp", "test.hpp"] + const parsers = await loadRequiredLanguageParsers(files, WASM_DIR) + expect(parsers.c).toBeDefined() + expect(parsers.h).toBeDefined() + expect(parsers.cpp).toBeDefined() + expect(parsers.hpp).toBeDefined() + }) - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-javascript.wasm")) - expect(parsers.js).toBeDefined() - expect(parsers.jsx).toBeDefined() - expect(parsers.js.query).toBeDefined() - expect(parsers.jsx.query).toBeDefined() - }) + it("should handle Kotlin files correctly", async () => { + const files = ["test.kt", "test.kts"] + const parsers = await loadRequiredLanguageParsers(files, WASM_DIR) + expect(parsers.kt).toBeDefined() + expect(parsers.kts).toBeDefined() + expect(parsers.kt.query).toBeDefined() + expect(parsers.kts.query).toBeDefined() + }) - it("should load TypeScript parser for .ts and .tsx files", async () => { - const files = ["test.ts", "test.tsx"] - const parsers = await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-typescript.wasm")) - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-tsx.wasm")) - expect(parsers.ts).toBeDefined() - expect(parsers.tsx).toBeDefined() - }) - - it("should load Python parser for .py files", async () => { - const files = ["test.py"] - const parsers = await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-python.wasm")) - expect(parsers.py).toBeDefined() - }) - - it("should load multiple language parsers as needed", async () => { - const files = ["test.js", "test.py", "test.rs", "test.go"] - const parsers = await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledTimes(4) - expect(parsers.js).toBeDefined() - expect(parsers.py).toBeDefined() - expect(parsers.rs).toBeDefined() - expect(parsers.go).toBeDefined() - }) - - it("should handle C/C++ files correctly", async () => { - const files = ["test.c", "test.h", "test.cpp", "test.hpp"] - const parsers = await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-c.wasm")) - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-cpp.wasm")) - expect(parsers.c).toBeDefined() - expect(parsers.h).toBeDefined() - expect(parsers.cpp).toBeDefined() - expect(parsers.hpp).toBeDefined() - }) - - it("should handle Kotlin files correctly", async () => { - const files = ["test.kt", "test.kts"] - const parsers = await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-kotlin.wasm")) - expect(parsers.kt).toBeDefined() - expect(parsers.kts).toBeDefined() - expect(parsers.kt.query).toBeDefined() - expect(parsers.kts.query).toBeDefined() - }) - - it("should throw error for unsupported file extensions", async () => { - const files = ["test.unsupported"] - - await expect(loadRequiredLanguageParsers(files)).rejects.toThrow("Unsupported language: unsupported") - }) - - it("should load each language only once for multiple files", async () => { - const files = ["test1.js", "test2.js", "test3.js"] - await loadRequiredLanguageParsers(files) - - expect(mockLanguageLoad).toHaveBeenCalledTimes(1) - expect(mockLanguageLoad).toHaveBeenCalledWith(expect.stringContaining("tree-sitter-javascript.wasm")) - }) - - it("should set language for each parser instance", async () => { - const files = ["test.js", "test.py"] - await loadRequiredLanguageParsers(files) - - expect(mockSetLanguage).toHaveBeenCalledTimes(2) - }) + it("should throw error for unsupported file extensions", async () => { + const files = ["test.unsupported"] + await expect(loadRequiredLanguageParsers(files, WASM_DIR)).rejects.toThrow("Unsupported language: unsupported") }) }) diff --git a/src/services/tree-sitter/languageParser.ts b/src/services/tree-sitter/languageParser.ts index 336f9919c0..a8ac0a9ead 100644 --- a/src/services/tree-sitter/languageParser.ts +++ b/src/services/tree-sitter/languageParser.ts @@ -1,5 +1,5 @@ import * as path from "path" -import { Parser, Query, Language } from "web-tree-sitter" +import { Parser as ParserT, Language as LanguageT, Query as QueryT } from "web-tree-sitter" import { javascriptQuery, typescriptQuery, @@ -32,30 +32,33 @@ import { export interface LanguageParser { [key: string]: { - parser: Parser - query: Query + parser: ParserT + query: QueryT } } -async function loadLanguage(langName: string) { - return await Language.load(path.join(__dirname, `tree-sitter-${langName}.wasm`)) +async function loadLanguage(langName: string, sourceDirectory?: string) { + const baseDir = sourceDirectory || __dirname + const wasmPath = path.join(baseDir, `tree-sitter-${langName}.wasm`) + + try { + const { Language } = require("web-tree-sitter") + return await Language.load(wasmPath) + } catch (error) { + console.error(`Error loading language: ${wasmPath}: ${error instanceof Error ? error.message : error}`) + throw error + } } let isParserInitialized = false -async function initializeParser() { - if (!isParserInitialized) { - await Parser.init() - isParserInitialized = true - } -} - /* Using node bindings for tree-sitter is problematic in vscode extensions because of incompatibility with electron. Going the .wasm route has the advantage of not having to build for multiple architectures. -We use web-tree-sitter and tree-sitter-wasms which provides auto-updating prebuilt WASM binaries for tree-sitter's language parsers. +We use web-tree-sitter and tree-sitter-wasms which provides auto-updating +prebuilt WASM binaries for tree-sitter's language parsers. This function loads WASM modules for relevant language parsers based on input files: 1. Extracts unique file extensions @@ -72,142 +75,157 @@ Sources: - https://github.com/tree-sitter/tree-sitter/blob/master/lib/binding_web/README.md - https://github.com/tree-sitter/tree-sitter/blob/master/lib/binding_web/test/query-test.js */ -export async function loadRequiredLanguageParsers(filesToParse: string[]): Promise { - await initializeParser() +export async function loadRequiredLanguageParsers(filesToParse: string[], sourceDirectory?: string) { + const { Parser, Query } = require("web-tree-sitter") + + if (!isParserInitialized) { + try { + await Parser.init() + isParserInitialized = true + } catch (error) { + console.error(`Error initializing parser: ${error instanceof Error ? error.message : error}`) + throw error + } + } + const extensionsToLoad = new Set(filesToParse.map((file) => path.extname(file).toLowerCase().slice(1))) const parsers: LanguageParser = {} + for (const ext of extensionsToLoad) { - let language: Language - let query: Query + let language: LanguageT + let query: QueryT let parserKey = ext // Default to using extension as key + switch (ext) { case "js": case "jsx": case "json": - language = await loadLanguage("javascript") - query = language.query(javascriptQuery) + language = await loadLanguage("javascript", sourceDirectory) + query = new Query(language, javascriptQuery) break case "ts": - language = await loadLanguage("typescript") - query = language.query(typescriptQuery) + language = await loadLanguage("typescript", sourceDirectory) + query = new Query(language, typescriptQuery) break case "tsx": - language = await loadLanguage("tsx") - query = language.query(tsxQuery) + language = await loadLanguage("tsx", sourceDirectory) + query = new Query(language, tsxQuery) break case "py": - language = await loadLanguage("python") - query = language.query(pythonQuery) + language = await loadLanguage("python", sourceDirectory) + query = new Query(language, pythonQuery) break case "rs": - language = await loadLanguage("rust") - query = language.query(rustQuery) + language = await loadLanguage("rust", sourceDirectory) + query = new Query(language, rustQuery) break case "go": - language = await loadLanguage("go") - query = language.query(goQuery) + language = await loadLanguage("go", sourceDirectory) + query = new Query(language, goQuery) break case "cpp": case "hpp": - language = await loadLanguage("cpp") - query = language.query(cppQuery) + language = await loadLanguage("cpp", sourceDirectory) + query = new Query(language, cppQuery) break case "c": case "h": - language = await loadLanguage("c") - query = language.query(cQuery) + language = await loadLanguage("c", sourceDirectory) + query = new Query(language, cQuery) break case "cs": - language = await loadLanguage("c_sharp") - query = language.query(csharpQuery) + language = await loadLanguage("c_sharp", sourceDirectory) + query = new Query(language, csharpQuery) break case "rb": - language = await loadLanguage("ruby") - query = language.query(rubyQuery) + language = await loadLanguage("ruby", sourceDirectory) + query = new Query(language, rubyQuery) break case "java": - language = await loadLanguage("java") - query = language.query(javaQuery) + language = await loadLanguage("java", sourceDirectory) + query = new Query(language, javaQuery) break case "php": - language = await loadLanguage("php") - query = language.query(phpQuery) + language = await loadLanguage("php", sourceDirectory) + query = new Query(language, phpQuery) break case "swift": - language = await loadLanguage("swift") - query = language.query(swiftQuery) + language = await loadLanguage("swift", sourceDirectory) + query = new Query(language, swiftQuery) break case "kt": case "kts": - language = await loadLanguage("kotlin") - query = language.query(kotlinQuery) + language = await loadLanguage("kotlin", sourceDirectory) + query = new Query(language, kotlinQuery) break case "css": - language = await loadLanguage("css") - query = language.query(cssQuery) + language = await loadLanguage("css", sourceDirectory) + query = new Query(language, cssQuery) break case "html": - language = await loadLanguage("html") - query = language.query(htmlQuery) + language = await loadLanguage("html", sourceDirectory) + query = new Query(language, htmlQuery) break case "ml": case "mli": - language = await loadLanguage("ocaml") - query = language.query(ocamlQuery) + language = await loadLanguage("ocaml", sourceDirectory) + query = new Query(language, ocamlQuery) break case "scala": - language = await loadLanguage("scala") - query = language.query(luaQuery) // Temporarily use Lua query until Scala is implemented + language = await loadLanguage("scala", sourceDirectory) + query = new Query(language, luaQuery) // Temporarily use Lua query until Scala is implemented break case "sol": - language = await loadLanguage("solidity") - query = language.query(solidityQuery) + language = await loadLanguage("solidity", sourceDirectory) + query = new Query(language, solidityQuery) break case "toml": - language = await loadLanguage("toml") - query = language.query(tomlQuery) + language = await loadLanguage("toml", sourceDirectory) + query = new Query(language, tomlQuery) break case "vue": - language = await loadLanguage("vue") - query = language.query(vueQuery) + language = await loadLanguage("vue", sourceDirectory) + query = new Query(language, vueQuery) break case "lua": - language = await loadLanguage("lua") - query = language.query(luaQuery) + language = await loadLanguage("lua", sourceDirectory) + query = new Query(language, luaQuery) break case "rdl": - language = await loadLanguage("systemrdl") - query = language.query(systemrdlQuery) + language = await loadLanguage("systemrdl", sourceDirectory) + query = new Query(language, systemrdlQuery) break case "tla": - language = await loadLanguage("tlaplus") - query = language.query(tlaPlusQuery) + language = await loadLanguage("tlaplus", sourceDirectory) + query = new Query(language, tlaPlusQuery) break case "zig": - language = await loadLanguage("zig") - query = language.query(zigQuery) + language = await loadLanguage("zig", sourceDirectory) + query = new Query(language, zigQuery) break case "ejs": case "erb": - language = await loadLanguage("embedded_template") - parserKey = "embedded_template" // Use same key for both extensions - query = language.query(embeddedTemplateQuery) + parserKey = "embedded_template" // Use same key for both extensions. + language = await loadLanguage("embedded_template", sourceDirectory) + query = new Query(language, embeddedTemplateQuery) break case "el": - language = await loadLanguage("elisp") - query = language.query(elispQuery) + language = await loadLanguage("elisp", sourceDirectory) + query = new Query(language, elispQuery) break case "ex": case "exs": - language = await loadLanguage("elixir") - query = language.query(elixirQuery) + language = await loadLanguage("elixir", sourceDirectory) + query = new Query(language, elixirQuery) break default: throw new Error(`Unsupported language: ${ext}`) } + const parser = new Parser() parser.setLanguage(language) parsers[parserKey] = { parser, query } } + return parsers }