Skip to content

Commit f0a2d4e

Browse files
authored
Merge pull request #1458 from codeflash-ai/codeflash/optimize-pr1457-2026-02-12T04.58.15
⚡️ Speed up function `extract_dependent_function` by 197% in PR #1457 (`fix-coverage-qualified-name`)
2 parents c4ed6e3 + 48817d7 commit f0a2d4e

1 file changed

Lines changed: 20 additions & 7 deletions

File tree

codeflash/code_utils/coverage_utils.py

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,29 @@
1313
def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]:
1414
"""Extract the single dependent function from the code context excluding the main function."""
1515
dependent_functions = set()
16-
for code_string in code_context.testgen_context.code_strings:
17-
ast_tree = ast.parse(code_string.code)
18-
dependent_functions.update(
19-
{node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))}
20-
)
2116

2217
# Compare using bare name since AST extracts bare function names
2318
bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function
24-
if bare_main in dependent_functions:
25-
dependent_functions.discard(bare_main)
19+
20+
for code_string in code_context.testgen_context.code_strings:
21+
# Quick heuristic: skip parsing entirely if there is no 'def' token,
22+
# since no function definitions can be present without it.
23+
if "def" not in code_string.code:
24+
continue
25+
26+
ast_tree = ast.parse(code_string.code)
27+
# Add function names directly, skipping the bare main name.
28+
for node in ast_tree.body:
29+
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
30+
name = node.name
31+
if name == bare_main:
32+
continue
33+
dependent_functions.add(name)
34+
# If more than one dependent function (other than the main) is found,
35+
# we can return False early since the final result cannot be a single name.
36+
if len(dependent_functions) > 1:
37+
return False
38+
2639

2740
if not dependent_functions:
2841
return False

0 commit comments

Comments
 (0)