Skip to content

⚡️ Speed up function _parse_and_collect_imports by 69% in PR #1498 (cf-simplify-context-extraction) #1499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
KRRT7 merged 4 commits into from
Feb 16, 2026
Merged

⚡️ Speed up function _parse_and_collect_imports by 69% in PR #1498 (cf-simplify-context-extraction) #1499

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bace611
Optimize _parse_and_collect_imports
codeflash-ai[bot] Feb 16, 2026
73e71d0
style: auto-fix linting issues
github-actions[bot] Feb 16, 2026
29c0a66
fix: resolve mypy type errors in collect_imports
github-actions[bot] Feb 16, 2026
bfa55cb
fix: handle ast.Match (Python 3.10+) in collect_imports traversal
github-actions[bot] Feb 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 42 additions & 6 deletions codeflash/languages/python/context/code_context_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -553,12 +553,48 @@ def _parse_and_collect_imports(code_context: CodeStringsMarkdown) -> tuple[ast.M
except SyntaxError:
return None
imported_names: dict[str, str] = {}
for node in ast.walk(tree):
if isinstance(node, ast.ImportFrom) and node.module:
for alias in node.names:
if alias.name != "*":
imported_name = alias.asname if alias.asname else alias.name
imported_names[imported_name] = node.module

# Directly iterate over the module body and nested structures instead of ast.walk
# This avoids traversing every single node in the tree
def collect_imports(nodes: list[ast.stmt]) -> None:
for node in nodes:
if isinstance(node, ast.ImportFrom) and node.module:
for alias in node.names:
if alias.name != "*":
imported_name = alias.asname if alias.asname else alias.name
imported_names[imported_name] = node.module
# Recursively check nested structures (function defs, class defs, if statements, etc.)
elif isinstance(
node,
(
ast.FunctionDef,
ast.AsyncFunctionDef,
ast.ClassDef,
ast.If,
ast.For,
ast.AsyncFor,
ast.While,
ast.With,
ast.AsyncWith,
ast.Try,
ast.ExceptHandler,
),
Comment on lines +567 to +581
Copy link
Copy Markdown
Contributor

@claude claude Bot Feb 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Missing ast.Match (Python 3.10+) node type

The collect_imports recursive traversal handles Try, If, For, While, With, etc., but does not handle ast.Match (match/case statements, introduced in Python 3.10). If someone writes:

match value:
    case 1:
        from module import something

The ImportFrom inside the match_case body won't be found by this optimized traversal, whereas the original ast.walk() would find it.

Since the project targets Python 3.9+, this is a low-probability edge case (imports inside match/case are very uncommon), but it's a correctness gap vs the original implementation.

Consider adding ast.Match and match_case to the isinstance check (guarded by a version check or hasattr).

):
if hasattr(node, "body"):
collect_imports(node.body)
if hasattr(node, "orelse"):
collect_imports(node.orelse)
if hasattr(node, "finalbody"):
collect_imports(node.finalbody)
if hasattr(node, "handlers"):
for handler in node.handlers:
collect_imports(handler.body)
# Handle match/case statements (Python 3.10+)
elif hasattr(ast, "Match") and isinstance(node, ast.Match):
for case in node.cases:
collect_imports(case.body)

collect_imports(tree.body)
return tree, imported_names


Expand Down
Loading