Add Python import tracking to code graph

api/analyzers/analyzer.py

@@ -143,3 +143,32 @@ def resolve_symbol(self, files: dict[Path, File], lsp: SyncLanguageServer, file_
 
         pass
 
+    @abstractmethod
+    def add_file_imports(self, file: File) -> None:
+        """
+        Add import statements to the file.
+
+        Args:
+            file (File): The file to add imports to.
+        """
+
+        pass
+
+    @abstractmethod
+    def resolve_import(self, files: dict[Path, File], lsp: SyncLanguageServer, file_path: Path, path: Path, import_node: Node) -> list[Entity]:
+        """
+        Resolve an import statement to entities.
+
+        Args:
+            files (dict[Path, File]): All files in the project.
+            lsp (SyncLanguageServer): The language server.
+            file_path (Path): The path to the file containing the import.
+            path (Path): The path to the project root.
+            import_node (Node): The import statement node.
+
+        Returns:
+            list[Entity]: List of resolved entities.
+        """
+
+        pass
+

api/analyzers/java/analyzer.py

@@ -132,3 +132,19 @@ def resolve_symbol(self, files: dict[Path, File], lsp: SyncLanguageServer, file_
             return self.resolve_method(files, lsp, file_path, path, symbol)
         else:
             raise ValueError(f"Unknown key {key}")
+
+    def add_file_imports(self, file: File) -> None:
+        """
+        Extract and add import statements from the file.
+        Java imports are not yet implemented.
+        """
+        # TODO: Implement Java import tracking
+        pass
+
+    def resolve_import(self, files: dict[Path, File], lsp: SyncLanguageServer, file_path: Path, path: Path, import_node: Node) -> list[Entity]:
+        """
+        Resolve an import statement to the entities it imports.
+        Java imports are not yet implemented.
+        """
+        # TODO: Implement Java import resolution
+        return []

api/analyzers/python/analyzer.py

@@ -121,3 +121,95 @@ def resolve_symbol(self, files: dict[Path, File], lsp: SyncLanguageServer, file_
             return self.resolve_method(files, lsp, file_path, path, symbol)
         else:
             raise ValueError(f"Unknown key {key}")
+
+    def add_file_imports(self, file: File) -> None:
+        """
+        Extract and add import statements from the file.
+
+        Supports:
+        - import module
+        - import module as alias
+        - from module import name
+        - from module import name1, name2
+        - from module import name as alias
+        """
+        try:
+            import warnings
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                # Query for both import types
+                import_query = self.language.query("""
+                    (import_statement) @import
+                    (import_from_statement) @import_from
+                """)
+
+            captures = import_query.captures(file.tree.root_node)
+
+            # Add all import statement nodes to the file
+            if 'import' in captures:
+                for import_node in captures['import']:
+                    file.add_import(import_node)
+
+            if 'import_from' in captures:
+                for import_node in captures['import_from']:
+                    file.add_import(import_node)
+        except Exception as e:
+            logger.debug(f"Failed to extract imports from {file.path}: {e}")
+
+    def resolve_import(self, files: dict[Path, File], lsp: SyncLanguageServer, file_path: Path, path: Path, import_node: Node) -> list[Entity]:
+        """
+        Resolve an import statement to the entities it imports.
+        """
+        res = []
+
+        try:
+            if import_node.type == 'import_statement':
+                # Handle "import module" or "import module as alias"
+                # Find all dotted_name and aliased_import nodes
+                for child in import_node.children:
+                    if child.type == 'dotted_name':
+                        # Try to resolve the module/name
+                        identifier = child.children[0] if child.child_count > 0 else child
+                        resolved = self.resolve_type(files, lsp, file_path, path, identifier)
+                        res.extend(resolved)
+                    elif child.type == 'aliased_import':
+                        # Get the actual name from aliased import (before 'as')
+                        if child.child_count > 0:
+                            actual_name = child.children[0]
+                            if actual_name.type == 'dotted_name' and actual_name.child_count > 0:
+                                identifier = actual_name.children[0]
+                            else:
+                                identifier = actual_name
+                            resolved = self.resolve_type(files, lsp, file_path, path, identifier)
+                            res.extend(resolved)
+
+            elif import_node.type == 'import_from_statement':
+                # Handle "from module import name1, name2"
+                # Find the 'import' keyword to know where imported names start
+                import_keyword_found = False
+                for child in import_node.children:
+                    if child.type == 'import':
+                        import_keyword_found = True
+                        continue
+
+                    # After 'import' keyword, dotted_name nodes are the imported names
+                    if import_keyword_found and child.type == 'dotted_name':
+                        # Try to resolve the imported name
+                        identifier = child.children[0] if child.child_count > 0 else child
+                        resolved = self.resolve_type(files, lsp, file_path, path, identifier)
+                        res.extend(resolved)
+                    elif import_keyword_found and child.type == 'aliased_import':
+                        # Handle "from module import name as alias"
+                        if child.child_count > 0:
+                            actual_name = child.children[0]
+                            if actual_name.type == 'dotted_name' and actual_name.child_count > 0:
+                                identifier = actual_name.children[0]
+                            else:
+                                identifier = actual_name
+                            resolved = self.resolve_type(files, lsp, file_path, path, identifier)
+                            res.extend(resolved)
+
+        except Exception as e:
+            logger.debug(f"Failed to resolve import: {e}")
+
+        return res

api/analyzers/source_analyzer.py

@@ -114,6 +114,10 @@ def first_pass(self, path: Path, files: list[Path], ignore: list[str], graph: Gr
             # Walk thought the AST
             graph.add_file(file)
             self.create_hierarchy(file, analyzer, graph)
+
+            # Extract import statements
+            if not analyzer.is_dependency(str(file_path)):
+                analyzer.add_file_imports(file)
 
     def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
         """
@@ -148,6 +152,8 @@ def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
             for i, file_path in enumerate(files):
                 file = self.files[file_path]
                 logging.info(f'Processing file ({i + 1}/{files_len}): {file_path}')
+
+                # Resolve entity symbols
                 for _, entity in file.entities.items():
                     entity.resolved_symbol(lambda key, symbol: analyzers[file_path.suffix].resolve_symbol(self.files, lsps[file_path.suffix], file_path, path, key, symbol))
                     for key, symbols in entity.symbols.items():
@@ -167,6 +173,13 @@ def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
                                 graph.connect_entities("RETURNS", entity.id, resolved_symbol.id)
                             elif key == "parameters":
                                 graph.connect_entities("PARAMETERS", entity.id, resolved_symbol.id)
+
+                # Resolve file imports
+                for import_node in file.imports:
+                    resolved_entities = analyzers[file_path.suffix].resolve_import(self.files, lsps[file_path.suffix], file_path, path, import_node)
+                    for resolved_entity in resolved_entities:
+                        file.add_resolved_import(resolved_entity)
+                        graph.connect_entities("IMPORTS", file.id, resolved_entity.id)
 
     def analyze_files(self, files: list[Path], path: Path, graph: Graph) -> None:
         self.first_pass(path, files, [], graph)

api/entities/file.py

@@ -21,10 +21,30 @@ def __init__(self, path: Path, tree: Tree) -> None:
         self.path = path
         self.tree = tree
         self.entities: dict[Node, Entity] = {}
+        self.imports: list[Node] = []
+        self.resolved_imports: set[Entity] = set()
 
     def add_entity(self, entity: Entity):
         entity.parent = self
         self.entities[entity.node] = entity
+
+    def add_import(self, import_node: Node):
+        """
+        Add an import statement node to track.
+
+        Args:
+            import_node (Node): The import statement node.
+        """
+        self.imports.append(import_node)
+
+    def add_resolved_import(self, resolved_entity: Entity):
+        """
+        Add a resolved import entity.
+
+        Args:
+            resolved_entity (Entity): The resolved entity that is imported.
+        """
+        self.resolved_imports.add(resolved_entity)
 
     def __str__(self) -> str:
         return f"path: {self.path}"

test-project/a.c

@@ -0,0 +1,11 @@
+#include <stdio.h>
+#include "/src/ff.h"
+
+
+/* Create an empty intset. */
+intset* intsetNew(void) {
+    intset *is = zmalloc(sizeof(intset));
+    is->encoding = intrev32ifbe(INTSET_ENC_INT16);
+    is->length = 0;
+    return is;
+}

test-project/b.py___

@@ -0,0 +1,73 @@
+class Employee:
+    def __init__(self, name, employee_id, base_salary):
+        self.name = name
+        self.employee_id = employee_id
+        self.base_salary = base_salary
+
+    def get_salary(self):
+        return self.base_salary
+
+    def display_info(self):
+        return f"Employee: {self.name} (ID: {self.employee_id})"
+
+    def apply_raise(self, percentage):
+        self.base_salary += (self.base_salary * percentage / 100)
+        return f"New salary after {percentage}% raise: ${self.base_salary:,.2f}"
+
+class Developer(Employee):
+    def __init__(self, name, employee_id, base_salary, programming_languages):
+        # Call parent class's __init__ method
+        super().__init__(name, employee_id, base_salary)
+        self.programming_languages = programming_languages
+
+    def add_language(self, language):
+        self.programming_languages.append(language)
+        return f"Added {language} to skill set"
+
+    # Override display_info method
+    def display_info(self):
+        basic_info = super().display_info()
+        return f"{basic_info}\nRole: Developer\nSkills: {', '.join(self.programming_languages)}"
+
+class Manager(Employee):
+    def __init__(self, name, employee_id, base_salary, team_members=None):
+        super().__init__(name, employee_id, base_salary)
+        self.team_members = team_members if team_members else []
+
+    def add_team_member(self, employee):
+        if isinstance(employee, Employee):
+            self.team_members.append(employee)
+            return f"Added {employee.name} to team"
+        return "Invalid team member"
+
+    def get_salary(self):
+        # Managers get 10% bonus on top of base salary
+        return self.base_salary * 1.1
+
+    def display_info(self):
+        basic_info = super().display_info()
+        team_names = [member.name for member in self.team_members]
+        return f"{basic_info}\nRole: Manager\nTeam Members: {', '.join(team_names)}"
+
+# Example usage
+if __name__ == "__main__":
+    # Create a developer
+    dev = Developer("Alice Smith", "DEV001", 75000, ["Python", "JavaScript"])
+    print(dev.display_info())
+    print(dev.add_language("Java"))
+    print(dev.apply_raise(10))
+
+    # Create another developer
+    dev2 = Developer("Bob Johnson", "DEV002", 70000, ["Python", "C++"]) 
+
+    # Create a manager
+    manager = Manager("Carol Williams", "MGR001", 100000)
+    print("\n" + manager.display_info())
+
+    # Add team members to manager
+    print(manager.add_team_member(dev))
+    print(manager.add_team_member(dev2))
+
+    # Display updated manager info
+    print("\n" + manager.display_info())
+    print(f"Manager's salary with bonus: ${manager.get_salary():,.2f}")

test-project/c.java

@@ -0,0 +1,26 @@
+package test-project;
+
+public class c {
+
+    private int a;
+
+    public static void main(String[] args) {
+        System.out.println("Hello, World!");
+    }
+
+    public static void print() {
+        System.out.println("Hello, World!");
+    }
+
+    public int getA() {
+        return a;
+    }   
+
+    public void setA(int a) {
+        this.a = a;
+    }
+
+    public void inc() {
+        setA(getA() + 1);
+    }
+}

tests/source_files/py_imports/module_a.py

@@ -0,0 +1,12 @@
+"""Module A with a class definition."""
+
+class ClassA:
+    """A simple class in module A."""
+
+    def method_a(self):
+        """A method in ClassA."""
+        return "Method A"
+
+def function_a():
+    """A function in module A."""
+    return "Function A"

tests/source_files/py_imports/module_b.py

@@ -0,0 +1,11 @@
+"""Module B that imports from module A."""
+
+from module_a import ClassA, function_a
+
+class ClassB(ClassA):
+    """A class that extends ClassA."""
+
+    def method_b(self):
+        """A method in ClassB."""
+        result = function_a()
+        return f"Method B: {result}"