feat: rich AST extraction with classes, methods, calls, and raises

- 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
This commit is contained in:
2026-06-10 15:58:52 +02:00
parent ed9f843d65
commit dbde651a2a
6 changed files with 589 additions and 55 deletions
+73 -13
View File
@@ -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
}
}
+4 -2
View File
@@ -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"
}
}
+433 -36
View File
@@ -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<string, string> = {
@@ -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") {
+12 -2
View File
@@ -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() };
+62 -1
View File
@@ -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<string>();
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])],
};
}
+5 -1
View File
@@ -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");
});