Merge pull request #869 from codeflash-ai/codeflash/optimize-pr868-2025-10-31T21.36.53

⚡️ Speed up method `ImportAnalyzer.visit_Attribute` by 38% in PR #868 (`import-analyser-fix`)

Author: misrasaurabh1

codeflash/discovery/discover_unit_tests.py

@@ -219,10 +219,18 @@ def __init__(self, function_names_to_find: set[str]) -> None:
  # For prefix match, store mapping from prefix-root to candidates for O(1) matching
  self._exact_names = function_names_to_find
  self._prefix_roots: dict[str, list[str]] = {}
+ # Precompute sets for faster lookup during visit_Attribute()
+ self._dot_names: set[str] = set()
+ self._dot_methods: dict[str, set[str]] = {}
+ self._class_method_to_target: dict[tuple[str, str], str] = {}
  for name in function_names_to_find:
  if "." in name:
- root = name.split(".", 1)[0]
- self._prefix_roots.setdefault(root, []).append(name)
+ root, method = name.rsplit(".", 1)
+ self._dot_names.add(name)
+ self._dot_methods.setdefault(method, set()).add(root)
+ self._class_method_to_target[(root, method)] = name
+ root_prefix = name.split(".", 1)[0]
+ self._prefix_roots.setdefault(root_prefix, []).append(name)
 
  def visit_Import(self, node: ast.Import) -> None:
  """Handle 'import module' statements."""
@@ -333,44 +341,45 @@ def visit_Attribute(self, node: ast.Attribute) -> None:
  if self.found_any_target_function:
  return
 
+ # Check if this is accessing a target function through an imported module
+
+ node_value = node.value
+ node_attr = node.attr
+
  # Check if this is accessing a target function through an imported module
  if (
- isinstance(node.value, ast.Name)
- and node.value.id in self.imported_modules
- and node.attr in self.function_names_to_find
+ isinstance(node_value, ast.Name)
+ and node_value.id in self.imported_modules
+ and node_attr in self.function_names_to_find
  ):
  self.found_any_target_function = True
- self.found_qualified_name = node.attr
+ self.found_qualified_name = node_attr
  return
 
- if isinstance(node.value, ast.Name) and node.value.id in self.imported_modules:
- for target_func in self.function_names_to_find:
- if "." in target_func:
- class_name, method_name = target_func.rsplit(".", 1)
- if node.attr == method_name:
- imported_name = node.value.id
- original_name = self.alias_mapping.get(imported_name, imported_name)
- if original_name == class_name:
- self.found_any_target_function = True
- self.found_qualified_name = target_func
- return
+ # Check for methods via imported modules using precomputed _dot_methods and _class_method_to_target
+ if isinstance(node_value, ast.Name) and node_value.id in self.imported_modules:
+ roots_possible = self._dot_methods.get(node_attr)
+ if roots_possible:
+ imported_name = node_value.id
+ original_name = self.alias_mapping.get(imported_name, imported_name)
+ if original_name in roots_possible:
+ self.found_any_target_function = True
+ self.found_qualified_name = self._class_method_to_target[(original_name, node_attr)]
+ 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
+ if isinstance(node_value, ast.Name) and node_value.id in self.instance_mapping:
+ class_name = self.instance_mapping[node_value.id]
+ roots_possible = self._dot_methods.get(node_attr)
+ if roots_possible and class_name in roots_possible:
+ self.found_any_target_function = True
+ self.found_qualified_name = self._class_method_to_target[(class_name, node_attr)]
+ 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:
+ # Check for dynamic import match
+ if self.has_dynamic_imports and node_attr in self.function_names_to_find:
  self.found_any_target_function = True
- self.found_qualified_name = node.attr
+ self.found_qualified_name = node_attr
  return
 
  self.generic_visit(node)