refactor: streamline import handling and enhance language query support

Author: SunBK201

scubatrace/cpp/file.py

@@ -1,34 +1,9 @@
 from functools import cached_property
 
 from ..file import File
-from . import language
 
 
 class CFile(File):
-    @cached_property
-    def imports(self) -> list[File]:
-        include_node = self.parser.query_all(
-            self.text, language.C.query_include
-        )
-        import_files = []
-        for node in include_node:
-            include_path_node = node.child_by_field_name("path")
-            if include_path_node is None:
-                continue
-            include = self.lsp.request_definition(
-                self.relpath,
-                include_path_node.start_point[0],
-                include_path_node.start_point[1],
-            )
-            if len(include) == 0:
-                continue
-            include = include[0]
-            include_abspath = include["absolutePath"]
-            if include_abspath not in self.project.files_abspath:
-                continue
-            import_files.append(self.project.files_abspath[include_abspath])
-        return import_files
-
     @cached_property
     def source_header(self) -> File | None:
         """

scubatrace/cpp/language.py

@@ -9,14 +9,22 @@ class C(Language):
     tslanguage = TSLanguage(tscpp.language())
 
     query_function = "(function_definition)@name"
-    query_identifier = "(identifier)@name"
     query_return = "(return_statement)@name"
     query_call = "(call_expression)@name"
+    query_import_identifier = """
+        (preproc_include
+            path: [
+                (system_lib_string)@name
+                (string_literal)@name
+            ]
+        )
+    """
+
     query_struct = "(struct_specifier)@name"
     query_class = "(class_specifier)@name"
-    query_method = "(function_definition)@name"
     query_field = "(field_declaration)@name"
     query_include = "(preproc_include)@name"
+
     query_global_statement = (
         "(declaration)@name"
         "(struct_specifier)@name"

scubatrace/file.py

@@ -192,7 +192,26 @@ def imports(self) -> list[File]:
         For example, in Python, this would include files imported using the `import` statement.
         In C/C++, this would include files included using the `#include` directive.
         """
-        ...
+        import_identifier_node = self.parser.query_all(
+            self.text, self.language.query_import_identifier
+        )
+        import_files = []
+        for node in import_identifier_node:
+            include = self.lsp.request_definition(
+                self.relpath,
+                node.start_point[0],
+                node.start_point[1],
+            )
+            if len(include) == 0:
+                continue
+            include = include[0]
+            include_abspath = include["absolutePath"]
+            if include_abspath in self.project.files_abspath:
+                import_files.append(self.project.files_abspath[include_abspath])
+            else:
+                # If the file is not in the project, we still add it to the imports
+                import_files.append(File.File(include_abspath, self.project))
+        return import_files
 
     @cached_property
     def functions(self) -> list[Function]:

scubatrace/java/file.py

@@ -1,5 +1,3 @@
-from functools import cached_property
-
 from ..file import File
 from . import language
 
@@ -12,34 +10,3 @@ def package(self) -> str:
             return ""
         package = package_node.text.decode()  # type: ignore
         return package
-
-    @cached_property
-    def imports(self) -> list[File]:
-        import_name_node = self.parser.query_all(
-            self.text, language.JAVA.query_import_name
-        )
-        import_files = []
-        for node in import_name_node:
-            include = self.lsp.request_definition(
-                self.relpath,
-                node.start_point[0],
-                node.start_point[1],
-            )
-            if len(include) == 0:
-                continue
-            include = include[0]
-            include_abspath = include["absolutePath"]
-            if include_abspath not in self.project.files_abspath:
-                continue
-            import_files.append(self.project.files_abspath[include_abspath])
-        return import_files
-
-    @cached_property
-    def import_class(self) -> list[str]:
-        import_node = self.parser.query_all(self.text, language.JAVA.query_import)
-        imports = []
-        for node in import_node:
-            assert node.text is not None
-            scoped_identifier = node.text.decode()
-            imports.append(scoped_identifier)
-        return imports

scubatrace/java/language.py

@@ -7,11 +7,23 @@
 class JAVA(Language):
     extensions = ["java"]
     tslanguage = TSLanguage(tsjava.language())
-    query_import = "(import_declaration(scoped_identifier)@name)"
+
+    query_function = query_method = "(method_declaration)@name"
+    query_return = "(return_statement)@name"
+    query_call = "(method_invocation)@name"
+    query_import_identifier = """
+        (import_declaration
+            (scoped_identifier
+                name: (identifier)@name
+            )
+        )
+        (import_declaration
+            (identifier)@name
+        )
+    """
+
     query_package = "(package_declaration)@name"
     query_class = "(class_declaration)@name"
-    query_function = query_method = "(method_declaration)@name"
-    query_identifier = "(identifier)@name"
 
     jump_statements = [
         "break_statement",
@@ -50,18 +62,12 @@ class JAVA(Language):
         "switch_expression",
     ]
 
-    loop_statements = ["for_statement", "while_statement", "do_statement", "enhanced_for_statement"]
-
-    query_import_name = """
-        (import_declaration
-            (scoped_identifier
-                name: (identifier)@name
-            )
-        )
-        (import_declaration
-            (identifier)@name
-        )
-    """
+    loop_statements = [
+        "for_statement",
+        "while_statement",
+        "do_statement",
+        "enhanced_for_statement",
+    ]
 
     @staticmethod
     def query_left_value(text):

scubatrace/javascript/language.py

@@ -8,12 +8,29 @@ class JAVASCRIPT(Language):
     extensions = ["js"]
     tslanguage = TSLanguage(tsjavascript.language())
 
-    query_function = "(function_declaration)@name"
-    query_method = "(method_definition)@name"
-    query_identifier = "(identifier)@name"
-    query_import = "(import_statement)@name"
-    query_export = "(export_statement)@name"
+    query_function = """
+    (function_declaration)@name
+    (method_definition)@name
+    """
+    query_return = "(return_statement)@name"
     query_call = "(call_expression)@name"
+    query_import = "(import_statement)@name"
+    query_import_identifier = """
+    (call_expression
+        function: [
+            (identifier)@require
+            (import)
+        ]
+        arguments: (arguments
+            (string)@name
+        )
+        (#eq? @require "require")
+    )
+    (import_statement
+        source: (string)@name
+    )
+    """
+
     query_class = "(class_declaration)@name"
 
     jump_statements = [

scubatrace/language.py

@@ -24,26 +24,33 @@ class Language:
     This is used to identify syntax errors in the code.
     """
 
-    query_function = "(function_definition)@name"
+    query_identifier = "(identifier)@name"
     """
-    The tree-sitter query to match function definitions.
+    The tree-sitter query to match identifiers.
     """
 
-    query_identifier = "(identifier)@name"
+    query_function: str
     """
-    The tree-sitter query to match identifiers.
+    The tree-sitter query to match function definitions.
     """
 
-    query_return = "(return_statement)@name"
+    query_return: str
     """
     The tree-sitter query to match return statements.
     """
 
-    query_call = "(call_expression)@name"
+    query_call: str
     """
     The tree-sitter query to match function calls.
     """
 
+    query_import_identifier: str
+    """
+    The tree-sitter query to match import identifiers.
+
+    For example, in C/C++, this would match the `header.h` in `#include <header.h>`.
+    """
+
     jump_statements: list[str]
     """
     The tree-sitter AST types of jump statements.

scubatrace/python/language.py

@@ -9,9 +9,28 @@ class PYTHON(Language):
     tslanguage = TSLanguage(tspython.language())
 
     query_function = "(function_definition)@name"
-    query_identifier = "(identifier)@name"
+    query_return = "(return_statement)@name"
+    query_call = "(call)@name"
+    query_import_identifier = """
+    (import_statement
+        name: [
+            (dotted_name)@name
+            (aliased_import
+                name: (dotted_name)@name
+            )
+        ]
+    )
+    (import_from_statement
+        module_name: [
+            (dotted_name)@name
+            (relative_import
+                (dotted_name)@name
+            )
+        ]
+    )
+    """
+
     query_class = "(class_definition)@name"
-    query_import = "(import_statement)@name"
 
     jump_statements = [
         "break_statement",

tests/c/test.c

@@ -1,4 +1,5 @@
 #include "include/sub.h"
+#include <stdio.h>
 int add(int a, int b)
 {
     return a + sub(a, b);
@@ -32,6 +33,7 @@ int main(int argc, char** argv)
         a += i;
         break;
     }
+    printf("a: %d, b: %d, count: %d\n", a, b, count);
     return 0;
 }
 

tests/javascript/test.js

@@ -1,7 +1,10 @@
 import "./admin.js";
 import { admin } from "./admin.js";
 var fs = require("fs");
-
+const fs = require("fs");
+const path = require("path");
+const module = await import('./admin.js');
+const { default: defaultExport } = await import('./admin.js');
 admin.name = "Pete";
 
 class Runoob {