From d8ec825d41b46224713fecb6cfc07a2fd3343a57 Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 19:13:10 +0000 Subject: [PATCH 1/5] Optimize collect_existing_class_names MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimization replaces `ast.walk(tree)` — which visits every node in the AST — with a manual stack-based traversal that only descends into container node types (`Module`, `ClassDef`, `FunctionDef`, control-flow statements, etc.) where `ClassDef` nodes can appear. This eliminates traversal of leaf nodes like `Name`, `Constant`, `Load`, and `Store`, which constitute the bulk of an AST but never contain class definitions. The profiler shows the original single-line comprehension spent 100% of runtime (117.7 ms) in `ast.walk`, while the optimized version completes in 36.1 ms (3.26× faster) by skipping ~60–80% of nodes depending on AST density. Tests confirm correctness across nested classes, control-flow scopes, and large trees with 1000+ classes. --- .../python/context/code_context_extractor.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 00db10e10..189e442a8 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -576,7 +576,25 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M def collect_existing_class_names(tree: ast.Module) -> set[str]: - return {node.name for node in ast.walk(tree) if isinstance(node, ast.ClassDef)} + class_names = set() + stack = [tree] + + while stack: + node = stack.pop() + + if isinstance(node, ast.ClassDef): + class_names.add(node.name) + + # Only traverse nodes that can contain ClassDef nodes + if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef, + ast.If, ast.For, ast.AsyncFor, ast.While, ast.With, ast.AsyncWith, + ast.Try, ast.ExceptHandler)): + stack.extend(getattr(node, 'body', [])) + stack.extend(getattr(node, 'orelse', [])) + stack.extend(getattr(node, 'finalbody', [])) + stack.extend(getattr(node, 'handlers', [])) + + return class_names BUILTIN_AND_TYPING_NAMES = frozenset( From dc9b41d90a58c4ea86d39f2110b3f14853b670d1 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 19:16:21 +0000 Subject: [PATCH 2/5] style: auto-fix ruff linting issues in collect_existing_class_names --- .../python/context/code_context_extractor.py | 36 +++++++++++++------ 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index 189e442a8..f75b0d9bc 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -578,22 +578,36 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M def collect_existing_class_names(tree: ast.Module) -> set[str]: class_names = set() stack = [tree] - + while stack: node = stack.pop() - + if isinstance(node, ast.ClassDef): class_names.add(node.name) - + # Only traverse nodes that can contain ClassDef nodes - if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef, - ast.If, ast.For, ast.AsyncFor, ast.While, ast.With, ast.AsyncWith, - ast.Try, ast.ExceptHandler)): - stack.extend(getattr(node, 'body', [])) - stack.extend(getattr(node, 'orelse', [])) - stack.extend(getattr(node, 'finalbody', [])) - stack.extend(getattr(node, 'handlers', [])) - + if isinstance( + node, + ( + ast.Module, + ast.ClassDef, + ast.FunctionDef, + ast.AsyncFunctionDef, + ast.If, + ast.For, + ast.AsyncFor, + ast.While, + ast.With, + ast.AsyncWith, + ast.Try, + ast.ExceptHandler, + ), + ): + stack.extend(getattr(node, "body", [])) + stack.extend(getattr(node, "orelse", [])) + stack.extend(getattr(node, "finalbody", [])) + stack.extend(getattr(node, "handlers", [])) + return class_names From 667a0a768d11ee39c14ca6fa891a713d0661b6a7 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 19:19:40 +0000 Subject: [PATCH 3/5] fix: resolve TC003 and mypy operator error in code_context_extractor --- codeflash/languages/python/context/code_context_extractor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index f75b0d9bc..d721844ec 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -6,7 +6,6 @@ from collections import defaultdict from dataclasses import dataclass, field from itertools import chain -from pathlib import Path from typing import TYPE_CHECKING import libcst as cst @@ -40,6 +39,8 @@ ) if TYPE_CHECKING: + from pathlib import Path + from jedi.api.classes import Name from codeflash.languages.base import DependencyResolver @@ -954,6 +955,7 @@ def _has_descriptor_like_class_fields(class_node: ast.ClassDef) -> bool: def _should_use_raw_project_class_context(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> bool: start_line = _get_class_start_line(class_node) + assert class_node.end_lineno is not None class_line_count = class_node.end_lineno - start_line + 1 is_small = ( class_line_count <= MAX_RAW_PROJECT_CLASS_LINES and len(class_node.body) <= MAX_RAW_PROJECT_CLASS_BODY_ITEMS From 2d5eec29499957e8a1710e19a43c1ea6618c2e6b Mon Sep 17 00:00:00 2001 From: "codeflash-ai[bot]" <148906541+codeflash-ai[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 19:35:14 +0000 Subject: [PATCH 4/5] Optimize _should_use_raw_project_class_context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The optimization replaced `any()` generator expressions with explicit early-return for-loops in four helper functions (`_is_namedtuple_class`, `_class_has_explicit_init`, `_has_descriptor_like_class_fields`, and `_has_non_property_method_decorator`), eliminating the overhead of building generator objects and calling the `any()` builtin. Line profiler data shows `_class_has_explicit_init` dropped from 1.85 ms to 0.96 ms (48% faster), and `_is_namedtuple_class` improved from 97 µs to 53 µs (46% faster), because the optimized code avoids allocating iterator state and returns immediately upon finding a match instead of completing the generator. The 51% overall runtime improvement (1.43 ms → 948 µs) comes from these cumulative reductions in per-call overhead across thousands of invocations during AST traversal. Test suite confirms no behavioral changes across all edge cases including dataclasses, decorators, and size-limit boundaries. --- .../python/context/code_context_extractor.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index d721844ec..f6d5f5a50 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -774,7 +774,10 @@ def _bool_literal(node: ast.AST) -> bool | None: def _is_namedtuple_class(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> bool: - return any(_expr_matches_name(base, import_aliases, "NamedTuple") for base in class_node.bases) + for base in class_node.bases: + if _expr_matches_name(base, import_aliases, "NamedTuple"): + return True + return False def _get_dataclass_config(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> tuple[bool, bool, bool]: @@ -812,10 +815,10 @@ def _get_class_start_line(class_node: ast.ClassDef) -> int: def _class_has_explicit_init(class_node: ast.ClassDef) -> bool: - return any( - isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == "__init__" - for item in class_node.body - ) + for item in class_node.body: + if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == "__init__": + return True + return False def _collect_synthetic_constructor_type_names(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> set[str]: @@ -948,9 +951,10 @@ def _has_non_property_method_decorator( def _has_descriptor_like_class_fields(class_node: ast.ClassDef) -> bool: - return any( - isinstance(item, (ast.Assign, ast.AnnAssign)) and isinstance(item.value, ast.Call) for item in class_node.body - ) + for item in class_node.body: + if isinstance(item, (ast.Assign, ast.AnnAssign)) and isinstance(item.value, ast.Call): + return True + return False def _should_use_raw_project_class_context(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> bool: From f9732eea6fbfb63d38d19f79f62f5d34a2305a18 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Mon, 16 Mar 2026 19:38:14 +0000 Subject: [PATCH 5/5] style: auto-fix ruff SIM110 linting issue in _is_namedtuple_class --- codeflash/languages/python/context/code_context_extractor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/languages/python/context/code_context_extractor.py b/codeflash/languages/python/context/code_context_extractor.py index f6d5f5a50..74b8d904b 100644 --- a/codeflash/languages/python/context/code_context_extractor.py +++ b/codeflash/languages/python/context/code_context_extractor.py @@ -774,7 +774,7 @@ def _bool_literal(node: ast.AST) -> bool | None: def _is_namedtuple_class(class_node: ast.ClassDef, import_aliases: dict[str, str]) -> bool: - for base in class_node.bases: + for base in class_node.bases: # noqa: SIM110 if _expr_matches_name(base, import_aliases, "NamedTuple"): return True return False