Improve import resolution logic and add test cases

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

Author: Copilot

api/analyzers/python/analyzer.py

@@ -152,42 +152,54 @@ def resolve_import(self, files: dict[Path, File], lsp: SyncLanguageServer, file_
         """
         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")
+        try:
             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)
+                # 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)
 
-        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)
+        except Exception as e:
+            logger.debug(f"Failed to resolve import: {e}")
 
         return res

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}"

tests/test_py_imports.py

@@ -0,0 +1,67 @@
+import os
+import unittest
+from pathlib import Path
+
+from api import SourceAnalyzer, File, Graph
+
+
+class Test_PY_Imports(unittest.TestCase):
+    def test_import_tracking(self):
+        """Test that Python imports are tracked correctly."""
+        # Get test file path
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        test_path = os.path.join(current_dir, 'source_files', 'py_imports')
+        
+        # Create graph and analyze
+        g = Graph("py_imports_test")
+        analyzer = SourceAnalyzer()
+        
+        try:
+            analyzer.analyze_local_folder(test_path, g)
+            
+            # Verify files were created
+            module_a = g.get_file('', 'module_a.py', '.py')
+            self.assertIsNotNone(module_a, "module_a.py should be in the graph")
+            
+            module_b = g.get_file('', 'module_b.py', '.py')
+            self.assertIsNotNone(module_b, "module_b.py should be in the graph")
+            
+            # Verify classes were created
+            class_a = g.get_class_by_name('ClassA')
+            self.assertIsNotNone(class_a, "ClassA should be in the graph")
+            
+            class_b = g.get_class_by_name('ClassB')
+            self.assertIsNotNone(class_b, "ClassB should be in the graph")
+            
+            # Verify function was created
+            func_a = g.get_function_by_name('function_a')
+            self.assertIsNotNone(func_a, "function_a should be in the graph")
+            
+            # Test: module_b should have IMPORTS relationship to ClassA
+            # Query to check if module_b imports ClassA
+            query = """
+                MATCH (f:File {name: 'module_b.py'})-[:IMPORTS]->(c:Class {name: 'ClassA'})
+                RETURN c
+            """
+            result = g._query(query, {})
+            self.assertGreater(len(result.result_set), 0, 
+                             "module_b.py should import ClassA")
+            
+            # Test: module_b should have IMPORTS relationship to function_a
+            query = """
+                MATCH (f:File {name: 'module_b.py'})-[:IMPORTS]->(fn:Function {name: 'function_a'})
+                RETURN fn
+            """
+            result = g._query(query, {})
+            self.assertGreater(len(result.result_set), 0, 
+                             "module_b.py should import function_a")
+            
+            print("✓ Import tracking test passed")
+            
+        finally:
+            # Cleanup: delete the test graph
+            g.delete()
+
+
+if __name__ == '__main__':
+    unittest.main()