using treesitter queries

Author: swilly22

app/api/repo/route.ts

@@ -1,40 +1,108 @@
-import { NextRequest, NextResponse } from "next/server";
-import { Graph, RedisClientType, RedisDefaultModules, createClient } from 'falkordb';
-import Parser from 'web-tree-sitter';
 import os from 'os';
 import fs from 'fs';
 import path from 'path';
 import simpleGit from 'simple-git';
+import Parser from 'web-tree-sitter';
+import { NextRequest, NextResponse } from "next/server";
+import { Graph, RedisClientType, RedisDefaultModules, createClient } from 'falkordb';
+
+// const client = createClient();
+// client.connect()
+
+//-----------------------------------------------------------------------------
+// Tree-Sitter queries
+//-----------------------------------------------------------------------------
+
+let parser;
+
+// class definition tree-sitter query
+// responsible for matching class definition, in addition to extracting the class name
+let class_definition_query;
+
+// function definition tree-sitter query
+// responsible for matching function definition, in addition to extracting the function name
+let function_definition_query;
+
+// function call tree-sitter query
+// responsible for matching function calls, in addition to extracting the callee function name
+let function_call_query;
+
+// Process Python file (Module)
+function processPythonSource(source_file) {
+	// read file
+	fs.readFile(source_file, 'utf8', function (err: any, source: string) {
+		if (err) {
+			return console.log(err);
+		}
+
+		// Construct an AST from source
+		const tree = parser.parse(source);
+		
+		// match all Class definitions
+		let class_matches = class_definition_query.matches(tree.rootNode);
+
+		// Iterate over each matched Class
+		for (let class_match of class_matches) {			
+			processClassDeclaration(source_file, class_match);
+		}
+	});
+}
+
+// Process Class declaration
+function processClassDeclaration(source_file, match) {
+	let class_node = match.captures[0].node;
+	let class_name = match.captures[1].node.text;
+	//console.log("Class node: " + class_node);
+	console.log("Class name: " + class_name);
+
+	// Match all function definition within the current class
+	let function_matches = function_definition_query.matches(class_node);
+	for (let function_match of function_matches) {	
+		processFunctionDeclaration(source_file, function_match);
+	}
+}
+
+// Process function declaration
+function processFunctionDeclaration(source_file, match) {
+	let function_node = match.captures[0].node;
+	let function_name = match.captures[1].node.text;
+	//console.log("Function definition Node: " + function_node);
+	console.log("Function name: " + function_name);
+
+	// Match all function calls within the current function
+	let function_call_matches = function_call_query.matches(function_node);
 
-//const client = createClient();
-//client.connect()
+	for (let function_call of function_call_matches) {
+		let callee = function_call.captures[0].node.text;
+		console.log(function_name + " calls: " + callee);
+	}
+}
 
 export async function POST(request: NextRequest) {
-	// parse source files
-	// const Parser = require('tree-sitter');
-	//await Parser.init();
+	// Initialize Tree-Sitter parser
 	await Parser.init({
 		locateFile(scriptName: string, scriptDirectory: string) {
 			return scriptName;
 		},
 	});
-	const parser = new Parser();
-	console.log("Parser: " + parser);
 
+	parser = new Parser();
 	const Python = await Parser.Language.load('tree-sitter-python.wasm');
-	console.log("Python: " + Python);
 
-	// const Python = require('tree-sitter-python');
-	//const JavaScript = require('tree-sitter-javascript');
-	// console.log("Python: " + Python);
-
-	//const parser = new Parser();
 	parser.setLanguage(Python);
 
+	//-------------------------------------------------------------------------
+	// Tree-Sitter AST queries
+	//-------------------------------------------------------------------------
+	
+	function_call_query = Python.query(`((call function: (identifier) @function-name))`);
+	//function_call_query       = Python.query(`(call_expression function: (identifier) @function-name)`);
+	class_definition_query    = Python.query(`(class_definition name: (identifier) @class-name) @class-definition`);
+	function_definition_query = Python.query(`(function_definition name: (identifier) @function-name) @function-definition`);
+
 	// Download Github Repo into a temporary folder
 	// Create temporary folder
 
-
 	const tmp_dir = os.tmpdir();
 	console.log("Temporary directory: " + tmp_dir);
 
@@ -71,29 +139,20 @@ export async function POST(request: NextRequest) {
 		for (let file of files) {
 			let file_path = path.join(dir, file);
 			let file_stat = fs.statSync(file_path);
-			if (file_stat.isDirectory()) {
-				walk(file_path);
-			}
-			else {
-				if (file.endsWith(".py")) {
-					source_files.push(file_path);
-				}
-			}
+			if (file_stat.isDirectory()) { walk(file_path); }
+			else if (file.endsWith(".py")) { source_files.push(file_path); }
 		}
 	}
 	walk(repo_root);
+	
+	//--------------------------------------------------------------------------
+	// process each source file
+	//--------------------------------------------------------------------------
 
 	for (let source_file of source_files) {
 		console.log("Processing file: " + source_file);
-		// read file
-		fs.readFile(source_file, 'utf8', function (err: any, source: string) {
-			if (err) {
-				return console.log(err);
-			}
-			const tree = parser.parse(source);
-			console.log("Tree: " + tree);
-		});
+		processPythonSource(source_file);
 	}
 
 	return NextResponse.json({ message: "in progress..." }, { status: 201 })
-}
+}