Merge branch 'refs/heads/main' into better-stdout-capture

misrasaurabh1 · misrasaurabh1 · commit 26f48b632ab8 · 2025-06-02T18:00:14.000-07:00
diff --git a/codeflash/api/cfapi.py b/codeflash/api/cfapi.py
@@ -115,7 +115,7 @@ def suggest_changes(
         "existingTests": existing_tests,
         "generatedTests": generated_tests,
         "traceId": trace_id,
-        "coverage": coverage_message,
+        "coverage_message": coverage_message,
     }
     return make_cfapi_request(endpoint="/suggest-pr-changes", method="POST", payload=payload)
 
@@ -151,7 +151,7 @@ def create_pr(
         "existingTests": existing_tests,
         "generatedTests": generated_tests,
         "traceId": trace_id,
-        "coverage": coverage_message,
+        "coverage_message": coverage_message,
     }
     return make_cfapi_request(endpoint="/create-pr", method="POST", payload=payload)
 
diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py
@@ -115,8 +115,7 @@ def create_trace_replay_test_code(
             if function_name == "__init__":
                 ret = {class_name_alias}(*args[1:], **kwargs)
             else:
-                instance = args[0] # self
-                ret = instance{method_name}(*args[1:], **kwargs)
+                ret = {class_name_alias}{method_name}(*args, **kwargs)
             """
     )
 
diff --git a/codeflash/code_utils/code_utils.py b/codeflash/code_utils/code_utils.py
@@ -15,6 +15,8 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.config_parser import find_pyproject_toml
 
+ImportErrorPattern = re.compile(r"ModuleNotFoundError.*$", re.MULTILINE)
+
 
 @contextmanager
 def custom_addopts() -> None:
diff --git a/codeflash/code_utils/config_consts.py b/codeflash/code_utils/config_consts.py
@@ -8,3 +8,4 @@
 N_TESTS_TO_GENERATE = 2
 TOTAL_LOOPING_TIME = 10.0  # 10 second candidate benchmarking budget
 COVERAGE_THRESHOLD = 60.0
+MIN_TESTCASE_PASSED_THRESHOLD = 6
diff --git a/codeflash/code_utils/env_utils.py b/codeflash/code_utils/env_utils.py
@@ -22,7 +22,7 @@ def check_formatter_installed(formatter_cmds: list[str], exit_on_failure: bool =
         f.flush()
         tmp_file = Path(f.name)
         try:
-            format_code(formatter_cmds, tmp_file)
+            format_code(formatter_cmds, tmp_file, print_status=False)
         except Exception:
             print(
                 "⚠️ Codeflash requires a code formatter to be installed in your environment, but none was found. Please install a supported formatter, verify the formatter-cmds in your codeflash pyproject.toml config and try again."
diff --git a/codeflash/code_utils/formatter.py b/codeflash/code_utils/formatter.py
@@ -13,7 +13,7 @@
     from pathlib import Path
 
 
-def format_code(formatter_cmds: list[str], path: Path) -> str:
+def format_code(formatter_cmds: list[str], path: Path, print_status: bool = True) -> str:  # noqa
     # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
     formatter_name = formatter_cmds[0].lower()
     if not path.exists():
@@ -28,7 +28,8 @@ def format_code(formatter_cmds: list[str], path: Path) -> str:
         try:
             result = subprocess.run(formatter_cmd_list, capture_output=True, check=False)
             if result.returncode == 0:
-                console.rule(f"Formatted Successfully with: {formatter_name.replace('$file', path.name)}")
+                if print_status:
+                    console.rule(f"Formatted Successfully with: {formatter_name.replace('$file', path.name)}")
             else:
                 logger.error(f"Failed to format code with {' '.join(formatter_cmd_list)}")
         except FileNotFoundError as e:
diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py
@@ -14,9 +14,16 @@
 
 import pytest
 from pydantic.dataclasses import dataclass
+from rich.panel import Panel
+from rich.text import Text
 
 from codeflash.cli_cmds.console import console, logger, test_files_progress_bar
-from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file, module_name_from_file_path
+from codeflash.code_utils.code_utils import (
+    ImportErrorPattern,
+    custom_addopts,
+    get_run_tmp_file,
+    module_name_from_file_path,
+)
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType
 
@@ -180,6 +187,10 @@ def discover_tests_pytest(
             logger.warning(
                 f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}\n {error_section}"
             )
+            if "ModuleNotFoundError" in result.stdout:
+                match = ImportErrorPattern.search(result.stdout).group()
+                panel = Panel(Text.from_markup(f"⚠️  {match} ", style="bold red"), expand=False)
+                console.print(panel)
 
         elif 0 <= exitcode <= 5:
             logger.warning(f"Failed to collect tests. Pytest Exit code: {exitcode}={pytest.ExitCode(exitcode).name}")
diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py
@@ -94,8 +94,6 @@ def visit_FunctionDef(self, node: FunctionDef) -> None:
             self.functions.append(
                 FunctionToOptimize(function_name=node.name, file_path=self.file_path, parents=self.ast_path[:])
             )
-        # Continue visiting the body of the function to find nested functions
-        self.generic_visit(node)
 
     def generic_visit(self, node: ast.AST) -> None:
         if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)):
diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py
@@ -23,6 +23,7 @@
 from codeflash.code_utils import env_utils
 from codeflash.code_utils.code_replacer import replace_function_definitions_in_module
 from codeflash.code_utils.code_utils import (
+    ImportErrorPattern,
     cleanup_paths,
     file_name_from_test_module_name,
     get_run_tmp_file,
@@ -1192,6 +1193,12 @@ def run_and_parse_tests(
                 f"stdout: {run_result.stdout}\n"
                 f"stderr: {run_result.stderr}\n"
             )
+            if "ModuleNotFoundError" in run_result.stdout:
+                from rich.text import Text
+
+                match = ImportErrorPattern.search(run_result.stdout).group()
+                panel = Panel(Text.from_markup(f"⚠️  {match} ", style="bold red"), expand=False)
+                console.print(panel)
         if testing_type in {TestingMode.BEHAVIOR, TestingMode.PERFORMANCE}:
             results, coverage_results = parse_test_results(
                 test_xml_path=result_file_path,
diff --git a/codeflash/result/critic.py b/codeflash/result/critic.py
@@ -4,7 +4,11 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils import env_utils
-from codeflash.code_utils.config_consts import COVERAGE_THRESHOLD, MIN_IMPROVEMENT_THRESHOLD
+from codeflash.code_utils.config_consts import (
+    COVERAGE_THRESHOLD,
+    MIN_IMPROVEMENT_THRESHOLD,
+    MIN_TESTCASE_PASSED_THRESHOLD,
+)
 from codeflash.models.models import TestType
 
 if TYPE_CHECKING:
@@ -50,7 +54,7 @@ def quantity_of_tests_critic(candidate_result: OptimizedCandidateResult) -> bool
     for test_type in report:
         pass_count += report[test_type]["passed"]
 
-    if pass_count >= 4:
+    if pass_count >= MIN_TESTCASE_PASSED_THRESHOLD:
         return True
     # If only one test passed, check if it's a REPLAY_TEST
     return bool(pass_count == 1 and report[TestType.REPLAY_TEST]["passed"] == 1)
diff --git a/tests/test_critic.py b/tests/test_critic.py
@@ -195,8 +195,7 @@ def test_generated_test_critic() -> None:
         timed_out=False,
         loop_index=1,
     )
-
-    test_results = [test_1, test_2, test_3, test_7]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_6, test_7, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -209,7 +208,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_6, test_7]
+    test_results = [test_1, test_2, test_3, test_6, test_7, test_1, test_4, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -222,7 +221,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_3, test_4, test_2, test_7]
+    test_results = [test_1, test_3, test_4, test_2, test_7, test_1, test_6, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -248,7 +247,7 @@ def test_generated_test_critic() -> None:
 
     assert not quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_4, test_5]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -287,7 +286,7 @@ def test_generated_test_critic() -> None:
 
     assert quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_4, test_5]
+    test_results = [test_1, test_2, test_3, test_4, test_5, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
@@ -328,7 +327,7 @@ def test_generated_test_critic() -> None:
 
     assert not quantity_of_tests_critic(candidate_result)
 
-    test_results = [test_1, test_2, test_3, test_5]
+    test_results = [test_1, test_2, test_3, test_5, test_1, test_1, test_1, test_1]
 
     candidate_result = OptimizedCandidateResult(
         max_loop_count=5,
diff --git a/tests/test_function_discovery.py b/tests/test_function_discovery.py
@@ -35,6 +35,17 @@ def test_function_eligible_for_optimization() -> None:
     assert len(functions_found[Path(f.name)]) == 0
 
 
+    # we want to trigger an error in the function discovery
+    function = """def test_invalid_code():"""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(function)
+        f.flush()
+        functions_found = find_all_functions_in_file(Path(f.name))
+    assert functions_found == {}
+
+
+
+
 def test_find_top_level_function_or_method():
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
         f.write(
@@ -82,6 +93,15 @@ def non_classmethod_function(cls, name):
         ).is_top_level
         # needed because this will be traced with a class_name being passed
 
+    # we want to write invalid code to ensure that the function discovery does not crash
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+            """def functionA():
+"""
+        )
+        f.flush()
+        path_obj_name = Path(f.name)
+        assert not inspect_top_level_functions_or_methods(path_obj_name, "functionA")
 
 def test_class_method_discovery():
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
@@ -152,6 +172,133 @@ def functionA():
             assert functions[file][0].function_name == "functionA"
 
 
+def test_nested_function():
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+import copy 
+
+def propagate_attributes(
+    nodes: dict[str, dict], edges: list[dict], source_node_id: str, attribute: str
+) -> dict[str, dict]:
+    modified_nodes = copy.deepcopy(nodes)
+
+    # Build an adjacency list for faster traversal
+    adjacency = {}
+    for edge in edges:
+        src = edge["source"]
+        tgt = edge["target"]
+        if src not in adjacency:
+            adjacency[src] = []
+        adjacency[src].append(tgt)
+
+    # Track visited nodes to avoid cycles
+    visited = set()
+
+    def traverse(node_id):
+        if node_id in visited:
+            return
+        visited.add(node_id)
+
+        # Propagate attribute from source node
+        if (
+            node_id != source_node_id
+            and source_node_id in modified_nodes
+            and attribute in modified_nodes[source_node_id]
+        ):
+            if node_id in modified_nodes:
+                modified_nodes[node_id][attribute] = modified_nodes[source_node_id][
+                    attribute
+                ]
+
+        # Continue propagation to neighbors
+        for neighbor in adjacency.get(node_id, []):
+            traverse(neighbor)
+
+    traverse(source_node_id)
+    return modified_nodes
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+def outer_function():
+    def inner_function():
+        pass
+
+    return inner_function
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(
+"""
+def outer_function():
+    def inner_function():
+        pass
+    
+    def another_inner_function():
+        pass
+    return inner_function, another_inner_function
+"""
+        )
+        f.flush()
+        test_config = TestConfig(
+            tests_root="tests", project_root_path=".", test_framework="pytest", tests_project_rootdir=Path()
+        )
+        path_obj_name = Path(f.name)
+        functions, functions_count = get_functions_to_optimize(
+            optimize_all=None,
+            replay_test=None,
+            file=path_obj_name,
+            test_cfg=test_config,
+            only_get_this_function=None,
+            ignore_paths=[Path("/bruh/")],
+            project_root=path_obj_name.parent,
+            module_root=path_obj_name.parent,
+        )
+
+        assert len(functions) == 1
+        assert functions_count == 1
+
+
 def test_filter_files_optimized():
     tests_root = Path("tests").resolve()
     module_root = Path().resolve()
diff --git a/tests/test_test_runner.py b/tests/test_test_runner.py
diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py

Original file line number	Diff line number	Diff line change
`@@ -115,8 +115,7 @@ def create_trace_replay_test_code(`
`115`	`115`	`if function_name == "__init__":`
`116`	`116`	`ret = {class_name_alias}(args[1:], *kwargs)`
`117`	`117`	`else:`
`118`		`- instance = args[0] # self`
`119`		`- ret = instance{method_name}(args[1:], *kwargs)`
	`118`	`+ ret = {class_name_alias}{method_name}(args, *kwargs)`
`120`	`119`	`"""`
`121`	`120`	`)`
`122`	`121`
Original file line number	Diff line number	Diff line change
`@@ -94,8 +94,6 @@ def visit_FunctionDef(self, node: FunctionDef) -> None:`
`94`	`94`	`self.functions.append(`
`95`	`95`	`FunctionToOptimize(function_name=node.name, file_path=self.file_path, parents=self.ast_path[:])`
`96`	`96`	`)`
`97`		`- # Continue visiting the body of the function to find nested functions`
`98`		`- self.generic_visit(node)`
`99`	`97`
`100`	`98`	`def generic_visit(self, node: ast.AST) -> None:`
`101`	`99`	`if isinstance(node, (FunctionDef, AsyncFunctionDef, ClassDef)):`