⚡️ Speed up function extract_dependent_function by 197% in PR #1457 (fix-coverage-qualified-name)

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