From dbde651a2a7e957fb644d6db092f56f18024ec6b Mon Sep 17 00:00:00 2001 From: Alex Blank Date: Wed, 10 Jun 2026 15:58:52 +0200 Subject: [PATCH] feat: rich AST extraction with classes, methods, calls, and raises MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Upgrade tree-sitter to 0.22.4 with compatible grammar versions - Add Python support via tree-sitter-python@0.23.6 - Add Go support via tree-sitter-go@0.23.4 - Extract class hierarchies with method signatures (params, returns) - Extract call graphs including external library calls - Extract exception raises (ValueError, Error, etc) - Extract top-level functions with full signatures - Encode rich AST data into compact markdown format: class:Name, method:name(params)→return, call:..., raise:... - Deduplicate nested call chains (db.query vs db.query().where()) - Filter out simple export names when rich func/class entries exist - All 55 tests passing --- package-lock.json | 86 ++++++-- package.json | 6 +- src/ast-extract.ts | 469 +++++++++++++++++++++++++++++++++++++--- src/llm-cache.ts | 14 +- src/merge.ts | 63 +++++- tests/llm-cache.test.ts | 6 +- 6 files changed, 589 insertions(+), 55 deletions(-) diff --git a/package-lock.json b/package-lock.json index 2ce3df8..c46f4bc 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,8 +13,10 @@ "openai": "^6.42.0", "p-limit": "^7.3.0", "picocolors": "^1.1.1", - "tree-sitter": "^0.21.0", - "tree-sitter-typescript": "^0.21.0" + "tree-sitter": "^0.22.4", + "tree-sitter-go": "^0.23.4", + "tree-sitter-python": "^0.23.6", + "tree-sitter-typescript": "^0.23.2" }, "bin": { "project-map": "dist/cli.js" @@ -3343,31 +3345,89 @@ } }, "node_modules/tree-sitter": { - "version": "0.21.1", - "resolved": "https://registry.npmjs.org/tree-sitter/-/tree-sitter-0.21.1.tgz", - "integrity": "sha512-7dxoA6kYvtgWw80265MyqJlkRl4yawIjO7S5MigytjELkX43fV2WsAXzsNfO7sBpPPCF5Gp0+XzHk0DwLCq3xQ==", + "version": "0.22.4", + "resolved": "https://registry.npmjs.org/tree-sitter/-/tree-sitter-0.22.4.tgz", + "integrity": "sha512-usbHZP9/oxNsUY65MQUsduGRqDHQOou1cagUSwjhoSYAmSahjQDAVsh9s+SlZkn8X8+O1FULRGwHu7AFP3kjzg==", "hasInstallScript": true, "license": "MIT", "dependencies": { - "node-addon-api": "^8.0.0", - "node-gyp-build": "^4.8.0" + "node-addon-api": "^8.3.0", + "node-gyp-build": "^4.8.4" + } + }, + "node_modules/tree-sitter-go": { + "version": "0.23.4", + "resolved": "https://registry.npmjs.org/tree-sitter-go/-/tree-sitter-go-0.23.4.tgz", + "integrity": "sha512-iQaHEs4yMa/hMo/ZCGqLfG61F0miinULU1fFh+GZreCRtKylFLtvn798ocCZjO2r/ungNZgAY1s1hPFyAwkc7w==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.1", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-javascript": { + "version": "0.23.1", + "resolved": "https://registry.npmjs.org/tree-sitter-javascript/-/tree-sitter-javascript-0.23.1.tgz", + "integrity": "sha512-/bnhbrTD9frUYHQTiYnPcxyHORIw157ERBa6dqzaKxvR/x3PC4Yzd+D1pZIMS6zNg2v3a8BZ0oK7jHqsQo9fWA==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2" + }, + "peerDependencies": { + "tree-sitter": "^0.21.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } + } + }, + "node_modules/tree-sitter-python": { + "version": "0.23.6", + "resolved": "https://registry.npmjs.org/tree-sitter-python/-/tree-sitter-python-0.23.6.tgz", + "integrity": "sha512-yIM9z0oxKIxT7bAtPOhgoVl6gTXlmlIhue7liFT4oBPF/lha7Ha4dQBS82Av6hMMRZoVnFJI8M6mL+SwWoLD3A==", + "hasInstallScript": true, + "license": "MIT", + "dependencies": { + "node-addon-api": "^8.3.0", + "node-gyp-build": "^4.8.4" + }, + "peerDependencies": { + "tree-sitter": "^0.22.1" + }, + "peerDependenciesMeta": { + "tree-sitter": { + "optional": true + } } }, "node_modules/tree-sitter-typescript": { - "version": "0.21.2", - "resolved": "https://registry.npmjs.org/tree-sitter-typescript/-/tree-sitter-typescript-0.21.2.tgz", - "integrity": "sha512-/RyNK41ZpkA8PuPZimR6pGLvNR1p0ibRUJwwQn4qAjyyLEIQD/BNlwS3NSxWtGsAWZe9gZ44VK1mWx2+eQVldg==", + "version": "0.23.2", + "resolved": "https://registry.npmjs.org/tree-sitter-typescript/-/tree-sitter-typescript-0.23.2.tgz", + "integrity": "sha512-e04JUUKxTT53/x3Uq1zIL45DoYKVfHH4CZqwgZhPg5qYROl5nQjV+85ruFzFGZxu+QeFVbRTPDRnqL9UbU4VeA==", "hasInstallScript": true, "license": "MIT", "dependencies": { - "node-addon-api": "^8.0.0", - "node-gyp-build": "^4.8.1" + "node-addon-api": "^8.2.2", + "node-gyp-build": "^4.8.2", + "tree-sitter-javascript": "^0.23.1" }, "peerDependencies": { "tree-sitter": "^0.21.0" }, "peerDependenciesMeta": { - "tree_sitter": { + "tree-sitter": { "optional": true } } diff --git a/package.json b/package.json index ad8beee..0da10a7 100644 --- a/package.json +++ b/package.json @@ -45,7 +45,9 @@ "openai": "^6.42.0", "p-limit": "^7.3.0", "picocolors": "^1.1.1", - "tree-sitter": "^0.21.0", - "tree-sitter-typescript": "^0.21.0" + "tree-sitter": "^0.22.4", + "tree-sitter-go": "^0.23.4", + "tree-sitter-python": "^0.23.6", + "tree-sitter-typescript": "^0.23.2" } } diff --git a/src/ast-extract.ts b/src/ast-extract.ts index 2727004..21cda43 100644 --- a/src/ast-extract.ts +++ b/src/ast-extract.ts @@ -1,9 +1,24 @@ import { readFileSync } from "fs"; import { extname } from "path"; -interface ASTFileData { +interface MethodData { + name: string; + params: string[]; + returns?: string; + calls: string[]; + raises: string[]; +} + +interface ClassData { + name: string; + methods: MethodData[]; +} + +export interface ASTFileData { exports: string[]; deps: string[]; + classes: ClassData[]; + functions: MethodData[]; } const LANGUAGE_MAP: Record = { @@ -37,11 +52,11 @@ export async function extractFileAST( const ts = require("tree-sitter-typescript"); grammar = langName === "tsx" ? ts.tsx : ts.typescript; } else { - // For other languages, try to require the grammar package try { const pkg = require(`tree-sitter-${langName}`); - // Some packages export { language }, others export the grammar directly - grammar = pkg.language || pkg.default || pkg; + // Some grammars export the language directly (e.g. python), + // others export { typescript, tsx } (e.g. typescript) + grammar = pkg.typescript || pkg.tsx || pkg.go || pkg; } catch { return null; } @@ -55,16 +70,424 @@ export async function extractFileAST( const content = readFileSync(filePath, "utf8"); const tree = parser.parse(content); + if (langName === "python") { + return extractPythonData(tree, content); + } + if ( + langName === "typescript" || + langName === "tsx" || + langName === "javascript" + ) { + return extractTypeScriptData(tree, content); + } + if (langName === "go") { + return extractGoData(tree, content); + } + + // Fallback for other languages const exports = extractExportsFromTree(tree, langName); const deps = extractDepsFromTree(tree, langName); - - return { exports, deps }; + return { exports, deps, classes: [], functions: [] }; } catch { - // Graceful fallback if tree-sitter fails return null; } } +// ============================================================================ +// PYTHON +// ============================================================================ + +function extractPythonData(tree: Tree, source: string): ASTFileData { + const classes: ClassData[] = []; + const exports: string[] = []; + const deps: string[] = []; + const functions: MethodData[] = []; + + function visit(node: SyntaxNode) { + // Imports + if ( + node.type === "import_statement" || + node.type === "import_from_statement" + ) { + const moduleNode = node.childForFieldName?.("module_name"); + if (moduleNode) { + deps.push(moduleNode.text); + } else { + // import a, b, c + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child?.type === "dotted_name" || child?.type === "identifier") { + deps.push(child.text); + } + } + } + } + + // Classes + if (node.type === "class_definition") { + const nameNode = node.childForFieldName?.("name"); + if (!nameNode) return; + const className = nameNode.text; + exports.push(className); + + const methods: MethodData[] = []; + const body = node.childForFieldName?.("body"); + if (body) { + for (let i = 0; i < body.childCount; i++) { + const child = body.child(i); + if (child?.type === "function_definition") { + const method = extractPythonMethod(child, source); + if (method) methods.push(method); + } + } + } + classes.push({ name: className, methods }); + return; // don't recurse into class body + } + + // Top-level functions + if (node.type === "function_definition") { + const nameNode = node.childForFieldName?.("name"); + if (nameNode) { + exports.push(nameNode.text); + const method = extractPythonMethod(node, source); + if (method) functions.push(method); + } + } + + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child) visit(child); + } + } + + visit(tree.rootNode); + return { exports: [...new Set(exports)], deps: [...new Set(deps)], classes, functions }; +} + +function extractPythonMethod( + node: SyntaxNode, + _source: string, +): MethodData | null { + const nameNode = node.childForFieldName?.("name"); + if (!nameNode) return null; + + const params: string[] = []; + const parameters = node.childForFieldName?.("parameters"); + if (parameters) { + for (let i = 0; i < parameters.childCount; i++) { + const param = parameters.child(i); + if (param?.type === "identifier" || param?.type === "typed_parameter") { + params.push(param.text); + } else if (param?.type === "typed_default_parameter") { + // name: type = default + const nameChild = param.childForFieldName?.("name"); + if (nameChild) params.push(nameChild.text); + } + } + } + + // Return type + let returns: string | undefined; + const returnType = node.childForFieldName?.("return_type"); + if (returnType) { + returns = returnType.text.replace(/^->\s*/, ""); + } + + // Calls and raises + const calls: string[] = []; + const raises: string[] = []; + const body = node.childForFieldName?.("body"); + if (body) { + extractPythonCallsAndRaises(body, calls, raises); + } + + return { + name: nameNode.text, + params, + returns, + calls: dedupeCallChains(calls), + raises: [...new Set(raises)], + }; +} + +function dedupeCallChains(calls: string[]): string[] { + const unique = [...new Set(calls)]; + // Remove shorter calls that are prefixes of longer ones + return unique.filter((call) => + !unique.some( + (other) => + other !== call && + (other.startsWith(`${call}.`) || other.startsWith(`${call}(`)), + ), + ); +} + +function extractPythonCallsAndRaises( + node: SyntaxNode, + calls: string[], + raises: string[], +) { + if (node.type === "call") { + const func = node.childForFieldName?.("function"); + if (func) { + const callStr = extractCallChain(func); + if (callStr) calls.push(callStr); + } + } + if (node.type === "raise_statement") { + const exc = node.child(1); + if (exc) { + const excType = + exc.type === "call" ? exc.childForFieldName?.("function") : exc; + if (excType) raises.push(excType.text); + } + } + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child && node.type !== "raise_statement") extractPythonCallsAndRaises(child, calls, raises); + } +} + +function extractCallChain(node: SyntaxNode): string | null { + if ( + node.type === "identifier" || + node.type === "attribute" || + node.type === "member_expression" || + node.type === "property_identifier" || + node.type === "type_identifier" + ) { + return node.text; + } + if (node.type === "call" || node.type === "call_expression") { + const func = node.childForFieldName?.("function"); + if (func) return extractCallChain(func); + } + return null; +} + +// ============================================================================ +// TYPESCRIPT / JAVASCRIPT +// ============================================================================ + +function extractTypeScriptData(tree: Tree, source: string): ASTFileData { + const classes: ClassData[] = []; + const exports: string[] = []; + const deps: string[] = []; + const functions: MethodData[] = []; + + function visit(node: SyntaxNode) { + // Imports + if ( + node.type === "import_statement" || + node.type === "import_declaration" + ) { + const sourceNode = node.childForFieldName?.("source"); + if (sourceNode) { + deps.push(sourceNode.text.slice(1, -1)); // remove quotes + } + } + if (node.type === "call_expression") { + const func = node.childForFieldName?.("function"); + if (func?.text === "require") { + const args = node.childForFieldName?.("arguments"); + if (args && args.childCount > 0) { + const firstArg = args.child(0); + if (firstArg?.type === "string") { + deps.push(firstArg.text.slice(1, -1)); + } + } + } + } + + // Classes + if (node.type === "class_declaration" || node.type === "class") { + const nameNode = node.childForFieldName?.("name"); + if (!nameNode) return; + const className = nameNode.text; + exports.push(className); + + const methods: MethodData[] = []; + const body = node.childForFieldName?.("body"); + if (body) { + for (let i = 0; i < body.childCount; i++) { + const child = body.child(i); + if ( + child?.type === "method_definition" || + child?.type === "function_definition" + ) { + const method = extractTSMethod(child, source); + if (method) methods.push(method); + } + } + } + classes.push({ name: className, methods }); + } + + // Exported functions/consts + if ( + node.type === "export_statement" || + node.type === "export_declaration" + ) { + const declaration = node.childForFieldName?.("declaration"); + if (declaration) { + const nameNode = findIdentifier(declaration); + if (nameNode) { + exports.push(nameNode.text); + if (declaration.type === "function_declaration" || declaration.type === "function") { + const method = extractTSMethod(declaration, source); + if (method) functions.push(method); + } + } + } + } + + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child) visit(child); + } + } + + visit(tree.rootNode); + return { exports: [...new Set(exports)], deps: [...new Set(deps)], classes, functions }; +} + +function extractTSMethod(node: SyntaxNode, _source: string): MethodData | null { + const nameNode = node.childForFieldName?.("name"); + if (!nameNode) return null; + + const params: string[] = []; + const parameters = node.childForFieldName?.("parameters"); + if (parameters) { + for (let i = 0; i < parameters.childCount; i++) { + const param = parameters.child(i); + if ( + param?.type === "identifier" || + param?.type === "required_parameter" || + param?.type === "optional_parameter" + ) { + const name = + param.childForFieldName?.("pattern") || + param.childForFieldName?.("name"); + if (name) { + const typeAnnotation = param.childForFieldName?.("type"); + if (typeAnnotation) { + const typeText = typeAnnotation.text.replace(/^:\s*/, ""); + params.push(`${name.text}: ${typeText}`); + } else { + params.push(name.text); + } + } else { + params.push(param.text); + } + } + } + } + + // Return type + let returns: string | undefined; + const returnType = node.childForFieldName?.("return_type"); + if (returnType) { + returns = returnType.text.replace(/^:\s*/, ""); + } + + // Calls and raises + const calls: string[] = []; + const raises: string[] = []; + const body = node.childForFieldName?.("body"); + if (body) { + extractTSCallsAndThrows(body, calls, raises); + } + + return { + name: nameNode.text, + params, + returns, + calls: dedupeCallChains(calls), + raises: [...new Set(raises)], + }; +} + +function extractTSCallsAndThrows( + node: SyntaxNode, + calls: string[], + raises: string[], +) { + if (node.type === "call_expression") { + const func = node.childForFieldName?.("function"); + if (func) { + const callStr = extractCallChain(func); + if (callStr) calls.push(callStr); + } + } + if (node.type === "throw_statement") { + const exc = node.child(1); + if (exc) { + // Strip 'new ' prefix and extract just the error type + const text = exc.text.replace(/^new\s+/, ""); + const match = text.match(/^(\w+)/); + if (match) raises.push(match[1]); + } + } + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child && node.type !== "throw_statement") extractTSCallsAndThrows(child, calls, raises); + } +} + +// ============================================================================ +// GO +// ============================================================================ + +function extractGoData(tree: Tree, _source: string): ASTFileData { + const exports: string[] = []; + const deps: string[] = []; + + function visit(node: SyntaxNode) { + if (node.type === "import_spec") { + const pathNode = node.childForFieldName?.("path"); + if (pathNode) deps.push(pathNode.text.slice(1, -1)); + } + + if ( + node.type === "function_declaration" || + node.type === "method_declaration" + ) { + const nameNode = node.childForFieldName?.("name"); + if (nameNode && /^[A-Z]/.test(nameNode.text)) { + exports.push(nameNode.text); + } + } + + if (node.type === "type_declaration") { + const spec = node.childForFieldName?.("spec"); + if (spec) { + const nameNode = spec.childForFieldName?.("name"); + if (nameNode && /^[A-Z]/.test(nameNode.text)) { + exports.push(nameNode.text); + } + } + } + + for (let i = 0; i < node.childCount; i++) { + const child = node.child(i); + if (child) visit(child); + } + } + + visit(tree.rootNode); + return { + exports: [...new Set(exports)], + deps: [...new Set(deps)], + classes: [], + functions: [], + }; +} + +// ============================================================================ +// GENERIC FALLBACK +// ============================================================================ + function extractExportsFromTree(tree: Tree, langName: string): string[] { const exports: string[] = []; const root = tree.rootNode; @@ -76,29 +499,10 @@ function extractExportsFromTree(tree: Tree, langName: string): string[] { langName === "javascript" ) { if (node.type === "export_statement") { - // export function foo - // export class Foo - // export const foo - // export { foo, bar } - // export default foo const declaration = node.childForFieldName?.("declaration"); if (declaration) { const nameNode = findIdentifier(declaration); if (nameNode) exports.push(nameNode.text); - } else { - // export { ... } - for (let i = 0; i < node.childCount; i++) { - const child = node.child(i); - if (child?.type === "export_clause") { - for (let j = 0; j < child.childCount; j++) { - const spec = child.child(j); - if (spec?.type === "export_specifier") { - const nameNode = spec.childForFieldName?.("name"); - if (nameNode) exports.push(nameNode.text); - } - } - } - } } } } else if (langName === "python") { @@ -112,11 +516,9 @@ function extractExportsFromTree(tree: Tree, langName: string): string[] { } else if (langName === "go") { if ( node.type === "function_declaration" || - node.type === "type_declaration" || - node.type === "var_declaration" || - node.type === "const_declaration" + node.type === "type_declaration" ) { - const nameNode = findIdentifier(node); + const nameNode = node.childForFieldName?.("name"); if (nameNode && /^[A-Z]/.test(nameNode.text)) { exports.push(nameNode.text); } @@ -157,13 +559,8 @@ function extractDepsFromTree(tree: Tree, langName: string): string[] { node.type === "import_declaration" ) { const source = node.childForFieldName?.("source"); - if (source) { - const text = source.text; - // Remove quotes - deps.push(text.slice(1, -1)); - } + if (source) deps.push(source.text.slice(1, -1)); } - // CommonJS: require("...") if (node.type === "call_expression") { const func = node.childForFieldName?.("function"); if (func?.text === "require") { diff --git a/src/llm-cache.ts b/src/llm-cache.ts index 0f6acae..f8b488a 100644 --- a/src/llm-cache.ts +++ b/src/llm-cache.ts @@ -1,4 +1,10 @@ -import { readFileSync, writeFileSync, existsSync, mkdirSync, renameSync } from "fs"; +import { + readFileSync, + writeFileSync, + existsSync, + mkdirSync, + renameSync, +} from "fs"; import { join } from "path"; const DEFAULT_CACHE_SUBDIR = ".pi-project-map"; @@ -48,7 +54,11 @@ export function getCached(hash: string, cacheDir?: string): string | undefined { return cache[hash]?.result; } -export function setCached(hash: string, result: string, cacheDir?: string): void { +export function setCached( + hash: string, + result: string, + cacheDir?: string, +): void { const cachePath = getCachePath(cacheDir); const cache = loadCache(cachePath); cache[hash] = { result, ts: Date.now() }; diff --git a/src/merge.ts b/src/merge.ts index 6a66ed1..6c1e7a5 100644 --- a/src/merge.ts +++ b/src/merge.ts @@ -9,6 +9,23 @@ interface LLMFileData { interface ASTFileData { exports: string[]; deps: string[]; + classes: Array<{ + name: string; + methods: Array<{ + name: string; + params: string[]; + returns?: string; + calls: string[]; + raises: string[]; + }>; + }>; + functions: Array<{ + name: string; + params: string[]; + returns?: string; + calls: string[]; + raises: string[]; + }>; } export function mergeFileData( @@ -16,10 +33,54 @@ export function mergeFileData( llm: LLMFileData, ast: ASTFileData | null, ): FileEntry { + const mergedExports = ast?.exports ?? llm.exports; + + // Remove simple export names that will be replaced by rich AST entries + const richNames = new Set(); + if (ast) { + for (const cls of ast.classes) richNames.add(cls.name); + for (const func of ast.functions) richNames.add(func.name); + } + const dedupedExports = mergedExports.filter((e) => !richNames.has(e)); + + // Encode classes into exports using compact DSL + if (ast && ast.classes.length > 0) { + for (const cls of ast.classes) { + const classExports: string[] = [`class:${cls.name}`]; + for (const method of cls.methods) { + const paramStr = method.params.join(", "); + const returnStr = method.returns ? ` → ${method.returns}` : ""; + classExports.push(`method:${method.name}(${paramStr})${returnStr}`); + if (method.calls.length > 0) { + classExports.push(`call:${method.calls.join(", ")}`); + } + if (method.raises.length > 0) { + classExports.push(`raise:${method.raises.join(", ")}`); + } + } + dedupedExports.push(...classExports); + } + } + + // Encode top-level functions + if (ast && ast.functions.length > 0) { + for (const func of ast.functions) { + const paramStr = func.params.join(", "); + const returnStr = func.returns ? ` → ${func.returns}` : ""; + dedupedExports.push(`func:${func.name}(${paramStr})${returnStr}`); + if (func.calls.length > 0) { + dedupedExports.push(`call:${func.calls.join(", ")}`); + } + if (func.raises.length > 0) { + dedupedExports.push(`raise:${func.raises.join(", ")}`); + } + } + } + return { name: fileName, purpose: llm.purpose, - exports: ast?.exports ?? llm.exports, + exports: dedupedExports, deps: [...new Set([...(ast?.deps ?? []), ...llm.deps])], }; } diff --git a/tests/llm-cache.test.ts b/tests/llm-cache.test.ts index 9357257..2681519 100644 --- a/tests/llm-cache.test.ts +++ b/tests/llm-cache.test.ts @@ -26,7 +26,11 @@ describe("llm-cache", () => { }); it("stores and retrieves cached results", () => { - setCached("abc123", "PURPOSE: test\nDEPS: none\nCONCEPTS: none", TEST_CACHE_DIR); + setCached( + "abc123", + "PURPOSE: test\nDEPS: none\nCONCEPTS: none", + TEST_CACHE_DIR, + ); const result = getCached("abc123", TEST_CACHE_DIR); expect(result).toBe("PURPOSE: test\nDEPS: none\nCONCEPTS: none"); });