new way

Author: Codeflash Bot

codeflash/discovery/discover_unit_tests.py

@@ -212,6 +212,8 @@ def __init__(self, function_names_to_find: set[str]) -> None:
         self.wildcard_modules: set[str] = set()
         # Track aliases: alias_name -> original_name
         self.alias_mapping: dict[str, str] = {}
+        # Track instances: variable_name -> class_name
+        self.instance_mapping: dict[str, str] = {}
 
         # Precompute function_names for prefix search
         # For prefix match, store mapping from prefix-root to candidates for O(1) matching
@@ -247,6 +249,24 @@ def visit_Import(self, node: ast.Import) -> None:
                     self.found_qualified_name = target_func
                     return
 
+    def visit_Assign(self, node: ast.Assign) -> None:
+        """Track variable assignments, especially class instantiations."""
+        if self.found_any_target_function:
+            return
+
+        # Check if the assignment is a class instantiation
+        if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
+            class_name = node.value.func.id
+            if class_name in self.imported_modules:
+                # Track all target variables as instances of the imported class
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        # Map the variable to the actual class name (handling aliases)
+                        original_class = self.alias_mapping.get(class_name, class_name)
+                        self.instance_mapping[target.id] = original_class
+
+        self.generic_visit(node)
+
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         """Handle 'from module import name' statements."""
         if self.found_any_target_function:
@@ -296,17 +316,6 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
                     self.found_qualified_name = target_func
                     return
 
-            # Check if any target function is a method of the imported class/module
-            # e.g., importing Graph and looking for Graph.topologicalSort
-            # TODO will pick up all tests which have the same class name (could be coming from a different file)
-            for target_func in fnames:
-                if "." in target_func:
-                    class_name, method_name = target_func.split(".", 1)
-                    if aname == class_name:
-                        self.found_any_target_function = True
-                        self.found_qualified_name = target_func
-                        return
-
     def visit_Attribute(self, node: ast.Attribute) -> None:
         """Handle attribute access like module.function_name."""
         if self.found_any_target_function:
@@ -334,6 +343,17 @@ def visit_Attribute(self, node: ast.Attribute) -> None:
                             self.found_qualified_name = target_func
                             return
 
+        # Check if this is accessing a method on an instance variable
+        if isinstance(node.value, ast.Name) and node.value.id in self.instance_mapping:
+            class_name = self.instance_mapping[node.value.id]
+            for target_func in self.function_names_to_find:
+                if "." in target_func:
+                    target_class, method_name = target_func.rsplit(".", 1)
+                    if node.attr == method_name and class_name == target_class:
+                        self.found_any_target_function = True
+                        self.found_qualified_name = target_func
+                        return
+
         # Check if this is accessing a target function through a dynamically imported module
         # Only if we've detected dynamic imports are being used
         if self.has_dynamic_imports and node.attr in self.function_names_to_find:

tests/test_unit_test_discovery.py

@@ -1310,6 +1310,30 @@ def test_target():
 
         assert should_process is True
 
+def test_analyze_imports_method():
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from code_to_optimize.topological_sort import Graph
+
+
+def test_topological_sort():
+    g = Graph(6)
+    g.addEdge(5, 2)
+    g.addEdge(5, 0)
+    g.addEdge(4, 0)
+    g.addEdge(4, 1)
+    g.addEdge(2, 3)
+    g.addEdge(3, 1)
+
+    assert g.topologicalSort()[0] == [5, 4, 2, 3, 1, 0]
+"""
+        test_file.write_text(test_content)
+
+        target_functions = {"Graph.topologicalSort"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+
+        assert should_process is True
 
 def test_analyze_imports_aliased_class_method_negative():
     with tempfile.TemporaryDirectory() as tmpdirname: