From 8a07c5e263faa92f8acce4c8f19b611f14dd8d0f Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 27 Mar 2026 14:27:52 -0500 Subject: [PATCH 1/3] 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. --- .../python/context/code_context_extractor.py | 5 ++- .../context/unused_definition_remover.py | 31 ++++++++++--------- .../python/static_analysis/code_extractor.py | 29 +++++++++++++++-- 3 files changed, 46 insertions(+), 19 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index e94cede2d..9a70de5fd 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/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) diff --git a/codeflash/languages/python/context/unused_definition_remover.py b/codeflash/languages/python/context/unused_definition_remover.py index aaa0435f8..575797b3d 100644 --- a/codeflash/languages/python/context/unused_definition_remover.py +++ b/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 diff --git a/codeflash/languages/python/static_analysis/code_extractor.py b/codeflash/languages/python/static_analysis/code_extractor.py index 454aeac9a..5426909d1 100644 --- a/codeflash/languages/python/static_analysis/code_extractor.py +++ b/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)) From f180c3f854a3d5255e44d93f1c16df8d8d00a903 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 27 Mar 2026 15:56:29 -0500 Subject: [PATCH 2/3] fix: use tuple syntax for isinstance check (Python 3.9 compat) --- codeflash/languages/python/static_analysis/code_extractor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/languages/python/static_analysis/code_extractor.py b/codeflash/languages/python/static_analysis/code_extractor.py index 5426909d1..3f146b50c 100644 --- a/codeflash/languages/python/static_analysis/code_extractor.py +++ b/codeflash/languages/python/static_analysis/code_extractor.py @@ -433,7 +433,7 @@ def _has_aliased_future_imports(module: cst.Module) -> bool: if ( isinstance(s, cst.ImportFrom) and s.module is not None - and isinstance(s.module, cst.Attribute | cst.Name) + and isinstance(s.module, (cst.Attribute, cst.Name)) and hasattr(s.module, "value") and s.module.value == "__future__" and isinstance(s.names, (list, tuple)) From 03e2fff2954d033b15bff13557013e07119fafbc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 27 Mar 2026 15:57:37 -0500 Subject: [PATCH 3/3] refactor: simplify _has_aliased_future_imports check cst.Attribute branch was dead code since __future__ imports always use a plain Name node. --- codeflash/languages/python/static_analysis/code_extractor.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/codeflash/languages/python/static_analysis/code_extractor.py b/codeflash/languages/python/static_analysis/code_extractor.py index 3f146b50c..9d937e55e 100644 --- a/codeflash/languages/python/static_analysis/code_extractor.py +++ b/codeflash/languages/python/static_analysis/code_extractor.py @@ -432,9 +432,7 @@ def _has_aliased_future_imports(module: cst.Module) -> bool: 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 isinstance(s.module, cst.Name) and s.module.value == "__future__" and isinstance(s.names, (list, tuple)) and any(name.asname is not None for name in s.names)