perf: optimize context extraction pipeline (~2x speedup)

Eliminate redundant CST traversals in code context extraction by caching
dependency data, skipping unnecessary transforms, and removing MetadataWrapper.

Author: KRRT7

codeflash/languages/python/context/code_context_extractor.py

@@ -395,7 +395,10 @@ def extract_all_contexts_from_files(
         except ValueError:
             relative_path = file_path
 
-        cleaned = remove_unused_definitions_by_function_names(original_module, hoh_names)
+        # Collect definitions + dependencies once (expensive CST traversal), reuse for mark pass
+        base_defs = collect_top_level_defs_with_dependencies(original_module)
+        hoh_defs = mark_defs_for_functions(base_defs, hoh_names)
+        cleaned = remove_unused_definitions_by_function_names(original_module, hoh_names, defs_with_usages=hoh_defs)
 
         # Pre-compute source imports once for this file
         src_gathered = gather_source_imports(original_module, file_path, project_root_path)

codeflash/languages/python/context/unused_definition_remover.py

@@ -165,8 +165,6 @@ def get_section_names(node: cst.CSTNode) -> list[str]:
 class DependencyCollector(cst.CSTVisitor):
     """Collects dependencies between definitions using the visitor pattern with depth tracking."""
 
-    METADATA_DEPENDENCIES = (cst.metadata.ParentNodeProvider,)
-
     def __init__(self, definitions: dict[str, UsageInfo]) -> None:
         super().__init__()
         self.definitions = definitions
@@ -179,6 +177,8 @@ def __init__(self, definitions: dict[str, UsageInfo]) -> None:
         # Track if we're processing a top-level variable
         self.processing_variable = False
         self.current_variable_names = set()
+        # Track Name nodes that are the .attr part of Attribute nodes (by id)
+        self.attr_name_ids: set[int] = set()
 
     def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
         function_name = node.name.value
@@ -281,6 +281,12 @@ def visit_AnnAssign(self, node: cst.AnnAssign) -> None:
         self.processing_variable = False
         self.current_variable_names.clear()
 
+    def visit_Attribute(self, node: cst.Attribute) -> None:
+        self.attr_name_ids.add(id(node.attr))
+
+    def leave_Attribute(self, original_node: cst.Attribute) -> None:
+        self.attr_name_ids.discard(id(original_node.attr))
+
     def visit_Name(self, node: cst.Name) -> None:
         name = node.value
 
@@ -296,15 +302,11 @@ def visit_Name(self, node: cst.Name) -> None:
             # Skip if this Name is the .attr part of an Attribute (e.g., 'x' in 'self.x')
             # We only want to track the base/value of attribute access, not the attribute name itself
             if self.class_depth > 0:
-                parent = self.get_metadata(cst.metadata.ParentNodeProvider, node)
-                if parent is not None and isinstance(parent, cst.Attribute):
-                    # Check if this Name is the .attr (property name), not the .value (base)
-                    # If it's the .attr, skip it - attribute names aren't references to definitions
-                    if parent.attr is node:
-                        return
-                    # If it's the .value (base), only skip if it's self/cls
-                    if name in ("self", "cls"):
-                        return
+                if id(node) in self.attr_name_ids:
+                    return
+                # If it's the .value (base), only skip if it's self/cls
+                if name in ("self", "cls"):
+                    return
             self.definitions[self.current_top_level_name].dependencies.add(name)
 
 
@@ -409,17 +411,16 @@ def remove_unused_definitions_recursively(
 
 
 def collect_top_level_defs_with_dependencies(code: Union[str, cst.Module]) -> dict[str, UsageInfo]:
-    """Collect all top level definitions and their inter-definition dependencies (expensive CST traversal).
+    """Collect all top level definitions and their inter-definition dependencies via CST traversal.
 
     Returns a definitions dict with dependencies populated but no usage marks set.
     This result can be reused across multiple mark_defs_for_functions calls to avoid
-    repeating the expensive MetadataWrapper + DependencyCollector traversal.
+    repeating the DependencyCollector traversal.
     """
     module = code if isinstance(code, cst.Module) else cst.parse_module(code)
     definitions = collect_top_level_definitions(module)
-    wrapper = cst.MetadataWrapper(module)
     dependency_collector = DependencyCollector(definitions)
-    wrapper.visit(dependency_collector)
+    module.visit(dependency_collector)
     return definitions
 
 

codeflash/languages/python/static_analysis/code_extractor.py

@@ -426,8 +426,31 @@ def leave_ImportFrom(
         return updated_node
 
 
+def _has_aliased_future_imports(module: cst.Module) -> bool:
+    for stmt in module.body:
+        if isinstance(stmt, cst.SimpleStatementLine):
+            for s in stmt.body:
+                if (
+                    isinstance(s, cst.ImportFrom)
+                    and s.module is not None
+                    and isinstance(s.module, cst.Attribute | cst.Name)
+                    and hasattr(s.module, "value")
+                    and s.module.value == "__future__"
+                    and isinstance(s.names, (list, tuple))
+                    and any(name.asname is not None for name in s.names)
+                ):
+                    return True
+    return False
+
+
+def _strip_future_aliases(module: cst.Module) -> cst.Module:
+    if _has_aliased_future_imports(module):
+        return module.visit(FutureAliasedImportTransformer())
+    return module
+
+
 def delete___future___aliased_imports(module_code: str) -> str:
-    return cst.parse_module(module_code).visit(FutureAliasedImportTransformer()).code
+    return _strip_future_aliases(cst.parse_module(module_code)).code
 
 
 def add_global_assignments(src_module_code: str, dst_module_code: str) -> str:
@@ -555,9 +578,9 @@ def gather_source_imports(
     src_module_and_package: ModuleNameAndPackage = calculate_module_and_package(project_root, src_path)
     try:
         if isinstance(src_module_code, cst.Module):
-            src_module = src_module_code.visit(FutureAliasedImportTransformer())
+            src_module = _strip_future_aliases(src_module_code)
         else:
-            src_module = cst.parse_module(src_module_code).visit(FutureAliasedImportTransformer())
+            src_module = _strip_future_aliases(cst.parse_module(src_module_code))
 
         has_module_level_imports = any(
             isinstance(s, (cst.Import, cst.ImportFrom))