diff --git a/codeflash/code_utils/coverage_utils.py b/codeflash/code_utils/coverage_utils.py index 083e63d9a..84e2a114f 100644 --- a/codeflash/code_utils/coverage_utils.py +++ b/codeflash/code_utils/coverage_utils.py @@ -13,16 +13,29 @@ def extract_dependent_function(main_function: str, code_context: CodeOptimizationContext) -> str | Literal[False]: """Extract the single dependent function from the code context excluding the main function.""" dependent_functions = set() - for code_string in code_context.testgen_context.code_strings: - ast_tree = ast.parse(code_string.code) - dependent_functions.update( - {node.name for node in ast_tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} - ) # Compare using bare name since AST extracts bare function names bare_main = main_function.rsplit(".", 1)[-1] if "." in main_function else main_function - if bare_main in dependent_functions: - dependent_functions.discard(bare_main) + + for code_string in code_context.testgen_context.code_strings: + # Quick heuristic: skip parsing entirely if there is no 'def' token, + # since no function definitions can be present without it. + if "def" not in code_string.code: + continue + + ast_tree = ast.parse(code_string.code) + # Add function names directly, skipping the bare main name. + for node in ast_tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + name = node.name + if name == bare_main: + continue + dependent_functions.add(name) + # If more than one dependent function (other than the main) is found, + # we can return False early since the final result cannot be a single name. + if len(dependent_functions) > 1: + return False + if not dependent_functions: return False