codeflash-ai/codeflash
1
0
Fork 0
Code Issues 18 Pull requests 33 Projects Releases Packages Wiki Activity Actions 4.6k

Merge #1343: Speed up _collect_numerical_imports by 159%

Browse source
This commit is contained in:
Kevin Turcios 2026-02-04 00:17:38 -05:00
commit 60bd77675d
1 changed files with 5 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
codeflash/code_utils/code_extractor.py
View file

@ -1429,7 +1429,9 @@ def _collect_numerical_imports(tree: ast.Module) -> tuple[set[str], set[str]]:
numerical_names: set[str] = set()
modules_used: set[str] = set()
for node in ast.walk(tree):
stack: list[ast.AST] = [tree]
while stack:
node = stack.pop()
if isinstance(node, ast.Import):
for alias in node.names:
# import numpy or import numpy as np
@ -1451,6 +1453,8 @@ def _collect_numerical_imports(tree: ast.Module) -> tuple[set[str], set[str]]:
name = alias.asname if alias.asname else alias.name
numerical_names.add(name)
modules_used.add(module_root)
else:
stack.extend(ast.iter_child_nodes(node))
return numerical_names, modules_used