652 lines
16 KiB
TypeScript
652 lines
16 KiB
TypeScript
import { readFileSync } from "fs";
|
|
import { extname } from "path";
|
|
|
|
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> = {
|
|
".ts": "typescript",
|
|
".tsx": "tsx",
|
|
".js": "javascript",
|
|
".jsx": "javascript",
|
|
".mjs": "javascript",
|
|
".py": "python",
|
|
".go": "go",
|
|
".rs": "rust",
|
|
".java": "java",
|
|
".c": "c",
|
|
".cpp": "cpp",
|
|
".h": "c",
|
|
".rb": "ruby",
|
|
};
|
|
|
|
export async function extractFileAST(
|
|
filePath: string,
|
|
): Promise<ASTFileData | null> {
|
|
const ext = extname(filePath).toLowerCase();
|
|
const langName = LANGUAGE_MAP[ext];
|
|
if (!langName) return null;
|
|
|
|
try {
|
|
const Parser = require("tree-sitter");
|
|
let grammar: unknown;
|
|
|
|
if (langName === "typescript" || langName === "tsx") {
|
|
const ts = require("tree-sitter-typescript");
|
|
grammar = langName === "tsx" ? ts.tsx : ts.typescript;
|
|
} else {
|
|
try {
|
|
const pkg = require(`tree-sitter-${langName}`);
|
|
// 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;
|
|
}
|
|
}
|
|
|
|
if (!grammar) return null;
|
|
|
|
const parser = new Parser();
|
|
parser.setLanguage(grammar);
|
|
|
|
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, classes: [], functions: [] };
|
|
} catch {
|
|
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.replace(/\s+/g, " ").trim();
|
|
}
|
|
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;
|
|
|
|
function visit(node: SyntaxNode) {
|
|
if (
|
|
langName === "typescript" ||
|
|
langName === "tsx" ||
|
|
langName === "javascript"
|
|
) {
|
|
if (node.type === "export_statement") {
|
|
const declaration = node.childForFieldName?.("declaration");
|
|
if (declaration) {
|
|
const nameNode = findIdentifier(declaration);
|
|
if (nameNode) exports.push(nameNode.text);
|
|
}
|
|
}
|
|
} else if (langName === "python") {
|
|
if (
|
|
node.type === "function_definition" ||
|
|
node.type === "class_definition"
|
|
) {
|
|
const nameNode = node.childForFieldName?.("name");
|
|
if (nameNode) exports.push(nameNode.text);
|
|
}
|
|
} else if (langName === "go") {
|
|
if (
|
|
node.type === "function_declaration" ||
|
|
node.type === "type_declaration"
|
|
) {
|
|
const nameNode = node.childForFieldName?.("name");
|
|
if (nameNode && /^[A-Z]/.test(nameNode.text)) {
|
|
exports.push(nameNode.text);
|
|
}
|
|
}
|
|
} else if (langName === "rust") {
|
|
if (
|
|
node.type === "function_item" ||
|
|
node.type === "struct_item" ||
|
|
node.type === "enum_item"
|
|
) {
|
|
const nameNode = node.childForFieldName?.("name");
|
|
if (nameNode) exports.push(nameNode.text);
|
|
}
|
|
}
|
|
|
|
for (let i = 0; i < node.childCount; i++) {
|
|
const child = node.child(i);
|
|
if (child) visit(child);
|
|
}
|
|
}
|
|
|
|
visit(root);
|
|
return [...new Set(exports)];
|
|
}
|
|
|
|
function extractDepsFromTree(tree: Tree, langName: string): string[] {
|
|
const deps: string[] = [];
|
|
const root = tree.rootNode;
|
|
|
|
function visit(node: SyntaxNode) {
|
|
if (
|
|
langName === "typescript" ||
|
|
langName === "tsx" ||
|
|
langName === "javascript"
|
|
) {
|
|
if (
|
|
node.type === "import_statement" ||
|
|
node.type === "import_declaration"
|
|
) {
|
|
const source = node.childForFieldName?.("source");
|
|
if (source) deps.push(source.text.slice(1, -1));
|
|
}
|
|
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));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} else if (langName === "python") {
|
|
if (
|
|
node.type === "import_statement" ||
|
|
node.type === "import_from_statement"
|
|
) {
|
|
const nameNode = node.childForFieldName?.("name");
|
|
if (nameNode) deps.push(nameNode.text);
|
|
}
|
|
} else if (langName === "go") {
|
|
if (node.type === "import_spec") {
|
|
const pathNode = node.childForFieldName?.("path");
|
|
if (pathNode) deps.push(pathNode.text.slice(1, -1));
|
|
}
|
|
} else if (langName === "rust") {
|
|
if (node.type === "use_declaration") {
|
|
const argument = node.childForFieldName?.("argument");
|
|
if (argument) deps.push(argument.text);
|
|
}
|
|
}
|
|
|
|
for (let i = 0; i < node.childCount; i++) {
|
|
const child = node.child(i);
|
|
if (child) visit(child);
|
|
}
|
|
}
|
|
|
|
visit(root);
|
|
return [...new Set(deps)];
|
|
}
|
|
|
|
function findIdentifier(node: SyntaxNode): SyntaxNode | null {
|
|
if (
|
|
node.type === "identifier" ||
|
|
node.type === "type_identifier" ||
|
|
node.type === "property_identifier"
|
|
) {
|
|
return node;
|
|
}
|
|
for (let i = 0; i < node.childCount; i++) {
|
|
const child = node.child(i);
|
|
if (child) {
|
|
const found = findIdentifier(child);
|
|
if (found) return found;
|
|
}
|
|
}
|
|
return null;
|
|
}
|
|
|
|
// Type definitions for tree-sitter nodes (simplified)
|
|
interface Tree {
|
|
rootNode: SyntaxNode;
|
|
}
|
|
|
|
interface SyntaxNode {
|
|
type: string;
|
|
text: string;
|
|
childCount: number;
|
|
childForFieldName?(name: string): SyntaxNode | null;
|
|
child(index: number): SyntaxNode | null;
|
|
}
|