refactor to early exit

Author: KRRT7

codeflash/discovery/discover_unit_tests.py

@@ -141,109 +141,132 @@ def close(self) -> None:
 
 
 class ImportAnalyzer(ast.NodeVisitor):
-    """AST-based analyzer to find all imports in a test file."""
+    """AST-based analyzer to check if any qualified names from function_names_to_find are imported or used in a test file."""
 
     def __init__(self, function_names_to_find: set[str]) -> None:
         self.function_names_to_find = function_names_to_find
-        self.imported_names: set[str] = set()
-        self.imported_modules: set[str] = set()
-        self.found_target_functions: set[str] = set()
-        self.qualified_names_called: set[str] = set()
+        self.found_any_target_function: bool = False
+        self.imported_modules: set[str] = set()  # Track imported modules for usage analysis
+        self.has_dynamic_imports: bool = False
+        self.wildcard_modules: set[str] = set()
 
     def visit_Import(self, node: ast.Import) -> None:
         """Handle 'import module' statements."""
+        if self.found_any_target_function:
+            return
+
         for alias in node.names:
             module_name = alias.asname if alias.asname else alias.name
             self.imported_modules.add(module_name)
-            self.imported_names.add(module_name)
-        self.generic_visit(node)
+
+            # Check for dynamic import modules
+            if alias.name == "importlib":
+                self.has_dynamic_imports = True
+
+            # Check if module itself is a target qualified name
+            if module_name in self.function_names_to_find:
+                self.found_any_target_function = True
+                return
+            # Check if any target qualified name starts with this module
+            for target_func in self.function_names_to_find:
+                if target_func.startswith(f"{module_name}."):
+                    self.found_any_target_function = True
+                    return
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         """Handle 'from module import name' statements."""
-        if node.module:
-            self.imported_modules.add(node.module)
+        if self.found_any_target_function:
+            return
+
+        if not node.module:
+            return
 
         for alias in node.names:
             if alias.name == "*":
-                continue
-            imported_name = alias.asname if alias.asname else alias.name
-            self.imported_names.add(imported_name)
-            if alias.name in self.function_names_to_find:
-                self.found_target_functions.add(alias.name)
-            # Check for qualified name matches
-            if node.module:
+                self.wildcard_modules.add(node.module)
+            else:
+                imported_name = alias.asname if alias.asname else alias.name
+                self.imported_modules.add(imported_name)
+
+                # Check for dynamic import functions
+                if node.module == "importlib" and alias.name == "import_module":
+                    self.has_dynamic_imports = True
+
+                # Check if imported name is a target qualified name
+                if alias.name in self.function_names_to_find:
+                    self.found_any_target_function = True
+                    return
+                # Check if module.name forms a target qualified name
                 qualified_name = f"{node.module}.{alias.name}"
                 if qualified_name in self.function_names_to_find:
-                    self.found_target_functions.add(qualified_name)
-        self.generic_visit(node)
+                    self.found_any_target_function = True
+                    return
 
-    def visit_Call(self, node: ast.Call) -> None:
-        """Handle dynamic imports like importlib.import_module() or __import__()."""
+    def visit_Attribute(self, node: ast.Attribute) -> None:
+        """Handle attribute access like module.function_name."""
+        if self.found_any_target_function:
+            return
+
+        # Check if this is accessing a target function through an imported module
         if (
-            isinstance(node.func, ast.Name)
-            and node.func.id == "__import__"
-            and node.args
-            and isinstance(node.args[0], ast.Constant)
-            and isinstance(node.args[0].value, str)
+            isinstance(node.value, ast.Name)
+            and node.value.id in self.imported_modules
+            and node.attr in self.function_names_to_find
         ):
-            # __import__("module_name")
-            self.imported_modules.add(node.args[0].value)
-        elif (
-            isinstance(node.func, ast.Attribute)
-            and isinstance(node.func.value, ast.Name)
-            and node.func.value.id == "importlib"
-            and node.func.attr == "import_module"
-            and node.args
-            and isinstance(node.args[0], ast.Constant)
-            and isinstance(node.args[0].value, str)
-        ):
-            # importlib.import_module("module_name")
-            self.imported_modules.add(node.args[0].value)
+            self.found_any_target_function = True
+            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:
+            self.found_any_target_function = True
+            return
+
         self.generic_visit(node)
 
     def visit_Name(self, node: ast.Name) -> None:
-        """Check if any name usage matches our target functions."""
-        if node.id in self.function_names_to_find:
-            self.found_target_functions.add(node.id)
-        self.generic_visit(node)
+        """Handle direct name usage like target_function()."""
+        if self.found_any_target_function:
+            return
 
-    def visit_Attribute(self, node: ast.Attribute) -> None:
-        """Handle module.function_name patterns."""
-        if node.attr in self.function_names_to_find:
-            self.found_target_functions.add(node.attr)
-            if isinstance(node.value, ast.Name):
-                qualified_name = f"{node.value.id}.{node.attr}"
-                self.qualified_names_called.add(qualified_name)
-        self.generic_visit(node)
+        # Check for __import__ usage
+        if node.id == "__import__":
+            self.has_dynamic_imports = True
 
+        if node.id in self.function_names_to_find:
+            self.found_any_target_function = True
+            return
 
-def analyze_imports_in_test_file(test_file_path: Path | str, target_functions: set[str]) -> bool:
-    """Analyze imports in a test file to determine if it might test any target functions.
+        # Check if this name could come from a wildcard import
+        for wildcard_module in self.wildcard_modules:
+            for target_func in self.function_names_to_find:
+                # Check if target_func is from this wildcard module and name matches
+                if target_func.startswith(f"{wildcard_module}.") and target_func.endswith(f".{node.id}"):
+                    self.found_any_target_function = True
+                    return
 
-    Args:
-        test_file_path: Path to the test file
-        target_functions: Set of function names we're looking for
+        self.generic_visit(node)
 
-    Returns:
-        bool: True if the test file should be processed (contains relevant imports), False otherwise
+    def generic_visit(self, node: ast.AST) -> None:
+        """Override generic_visit to stop traversal if a target function is found."""
+        if self.found_any_target_function:
+            return
+        super().generic_visit(node)
 
-    """
-    if isinstance(test_file_path, str):
-        test_file_path = Path(test_file_path)
 
+def analyze_imports_in_test_file(test_file_path: Path | str, target_functions: set[str]) -> bool:
+    """Analyze a test file to see if it imports any of the target functions."""
     try:
-        with test_file_path.open("r", encoding="utf-8") as f:
-            content = f.read()
-
-        tree = ast.parse(content, filename=str(test_file_path))
+        with Path(test_file_path).open("r", encoding="utf-8") as f:
+            source_code = f.read()
+        tree = ast.parse(source_code, filename=str(test_file_path))
         analyzer = ImportAnalyzer(target_functions)
         analyzer.visit(tree)
-
-        return bool(analyzer.found_target_functions)
-
-    except (SyntaxError, UnicodeDecodeError, OSError) as e:
+    except (SyntaxError, FileNotFoundError) as e:
         logger.debug(f"Failed to analyze imports in {test_file_path}: {e}")
         return True
+    else:
+        return analyzer.found_any_target_function
 
 
 def filter_test_files_by_imports(

tests/test_unit_test_discovery.py

@@ -834,6 +834,96 @@ def test_something():
         assert should_process is False
 
 
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from mymodule import *
+
+def test_target():
+    assert target_function() is True
+"""
+        test_file.group
+        test_file.write_text(test_content)
+        
+        target_functions = {"mymodule.target_function"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+        
+        assert should_process is True
+
+
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from mymodule import *
+
+def test_target():
+    assert target_function_extended() is True
+"""
+        test_file.write_text(test_content)
+        
+        # Should not match - target_function != target_function_extended  
+        target_functions = {"mymodule.target_function"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+        
+        assert should_process is False
+
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from mymodule import *
+
+def test_something():
+    x = 42
+    assert x == 42
+"""
+        test_file.write_text(test_content)
+        
+        target_functions = {"mymodule.target_function"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+        
+        assert should_process is False
+
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from mymodule import *
+
+def test_something():
+    message = "calling target_function"
+    assert "target_function" in message
+"""
+        test_file.write_text(test_content)
+        
+        target_functions = {"mymodule.target_function"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+        
+        # String literals are ast.Constant nodes, not ast.Name nodes, so they don't match
+        assert should_process is False
+
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        test_file = Path(tmpdirname) / "test_example.py"
+        test_content = """
+from mymodule import target_function
+from othermodule import *
+
+def test_target():
+    assert target_function() is True
+    assert other_func() is True
+"""
+        test_file.write_text(test_content)
+        
+        target_functions = {"mymodule.target_function", "othermodule.other_func"}
+        should_process = analyze_imports_in_test_file(test_file, target_functions)
+        
+        assert should_process is True
+
+
+
+
 def test_analyze_imports_module_import():
     """Test module imports with function access patterns."""
     with tempfile.TemporaryDirectory() as tmpdirname: