Add Import support for Python - initial implementation

Co-authored-by: gkorland <753206+gkorland@users.noreply.github.com>

Author: Copilot

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,73 @@ 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.
+        """
+        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)
+
+    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 = []
+        
+        # For import statements like "import os" or "from pathlib import Path"
+        # We need to find the dotted_name nodes that represent the imported modules/names
+        import warnings
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            if import_node.type == 'import_statement':
+                # Handle "import module" or "import module as alias"
+                # Look for dotted_name or aliased_import
+                query = self.language.query("(dotted_name) @module (aliased_import) @aliased")
+            else:  # import_from_statement
+                # Handle "from module import name"
+                # Get the imported names (after the 'import' keyword)
+                query = self.language.query("""
+                    (import_from_statement
+                        (dotted_name) @imported_name)
+                """)
+        
+        captures = query.captures(import_node)
+        
+        # Try to resolve each imported name
+        if 'module' in captures:
+            for module_node in captures['module']:
+                resolved = self.resolve_type(files, lsp, file_path, path, module_node)
+                res.extend(resolved)
+        
+        if 'aliased' in captures:
+            for aliased_node in captures['aliased']:
+                # Get the actual module name from the aliased import
+                if aliased_node.child_count > 0:
+                    module_name_node = aliased_node.children[0]
+                    resolved = self.resolve_type(files, lsp, file_path, path, module_name_node)
+                    res.extend(resolved)
+        
+        if 'imported_name' in captures:
+            for name_node in captures['imported_name']:
+                resolved = self.resolve_type(files, lsp, file_path, path, name_node)
+                res.extend(resolved)
+        
+        return res

api/analyzers/source_analyzer.py

@@ -112,6 +112,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:
         """
@@ -141,6 +145,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.resolved_symbols.items():
@@ -157,6 +163,13 @@ def second_pass(self, graph: Graph, files: list[Path], path: Path) -> None:
                                 graph.connect_entities("RETURNS", entity.id, symbol.id)
                             elif key == "parameters":
                                 graph.connect_entities("PARAMETERS", entity.id, 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

@@ -1,5 +1,6 @@
 from pathlib import Path
 from tree_sitter import Node, Tree
+from typing import Self
 
 from api.entities.entity import Entity
 
@@ -21,10 +22,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[Self] = 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: Self):
+        """
+        Add a resolved import entity.
+        
+        Args:
+            resolved_entity (Self): The resolved entity that is imported.
+        """
+        self.resolved_imports.add(resolved_entity)
 
     def __str__(self) -> str:
         return f"path: {self.path}"