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..9d937e55e 100644 --- a/codeflash/languages/python/static_analysis/code_extractor.py +++ b/codeflash/languages/python/static_analysis/code_extractor.py @@ -426,8 +426,29 @@ 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 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) + ): + 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 +576,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))