Merge branch 'main' into lsp/config-suggestions-configured-values

Author: aseembits93

codeflash/code_utils/formatter.py

@@ -96,6 +96,19 @@ def is_diff_line(line: str) -> bool:
     return len(diff_lines)
 
 
+def format_generated_code(generated_test_source: str, formatter_cmds: list[str]) -> str:
+    with tempfile.TemporaryDirectory() as test_dir_str:
+        # try running formatter, if nothing changes (could be due to formatting failing or no actual formatting needed) return code with 2 or more newlines substituted with 2 newlines
+        original_temp = Path(test_dir_str) / "original_temp.py"
+        original_temp.write_text(generated_test_source, encoding="utf8")
+        _, formatted_code, changed = apply_formatter_cmds(
+            formatter_cmds, original_temp, test_dir_str, print_status=False, exit_on_failure=False
+        )
+        if not changed:
+            return re.sub(r"\n{2,}", "\n\n", formatted_code)
+    return formatted_code
+
+
 def format_code(
     formatter_cmds: list[str],
     path: Union[str, Path],
@@ -120,7 +133,7 @@ def format_code(
         original_code_lines = len(original_code.split("\n"))
 
         if check_diff and original_code_lines > 50:
-            # we dont' count the formatting diff for the optimized function as it should be well-formatted
+            # we don't count the formatting diff for the optimized function as it should be well-formatted
             original_code_without_opfunc = original_code.replace(optimized_code, "")
 
             original_temp = Path(test_dir_str) / "original_temp.py"

codeflash/models/models.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections import defaultdict
+from collections import Counter, defaultdict
 from typing import TYPE_CHECKING
 
 from rich.tree import Tree
@@ -675,6 +675,16 @@ def total_passed_runtime(self) -> int:
             [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()]
         )
 
+    def file_to_no_of_tests(self, test_functions_to_remove: list[str]) -> Counter[Path]:
+        map_gen_test_file_to_no_of_tests = Counter()
+        for gen_test_result in self.test_results:
+            if (
+                gen_test_result.test_type == TestType.GENERATED_REGRESSION
+                and gen_test_result.id.test_function_name not in test_functions_to_remove
+            ):
+                map_gen_test_file_to_no_of_tests[gen_test_result.file_name] += 1
+        return map_gen_test_file_to_no_of_tests
+
     def __iter__(self) -> Iterator[FunctionTestInvocation]:
         return iter(self.test_results)
 

codeflash/optimization/function_optimizer.py

@@ -55,7 +55,7 @@
     remove_functions_from_generated_tests,
 )
 from codeflash.code_utils.env_utils import get_pr_number
-from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.formatter import format_code, format_generated_code, sort_imports
 from codeflash.code_utils.git_utils import git_root_dir
 from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
 from codeflash.code_utils.line_profile_utils import add_decorator_imports
@@ -1401,6 +1401,9 @@ def process_review(
         generated_tests = remove_functions_from_generated_tests(
             generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove
         )
+        map_gen_test_file_to_no_of_tests = original_code_baseline.behavior_test_results.file_to_no_of_tests(
+            test_functions_to_remove
+        )
 
         original_runtime_by_test = original_code_baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
         optimized_runtime_by_test = (
@@ -1413,11 +1416,16 @@ def process_review(
 
         generated_tests_str = ""
         for test in generated_tests.generated_tests:
-            generated_tests_str += f"```python\n{test.generated_original_test_source}\n```"
-            generated_tests_str += "\n\n"
+            if map_gen_test_file_to_no_of_tests[test.behavior_file_path] > 0:
+                formatted_generated_test = format_generated_code(
+                    test.generated_original_test_source, self.args.formatter_cmds
+                )
+                generated_tests_str += f"```python\n{formatted_generated_test}\n```"
+                generated_tests_str += "\n\n"
 
         if concolic_test_str:
-            generated_tests_str += f"```python\n{concolic_test_str}\n```\n\n"
+            formatted_generated_test = format_generated_code(concolic_test_str, self.args.formatter_cmds)
+            generated_tests_str += f"```python\n{formatted_generated_test}\n```\n\n"
 
         existing_tests, replay_tests, concolic_tests = existing_tests_source_for(
             self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root),
@@ -1533,7 +1541,7 @@ def process_review(
                 trace_id=self.function_trace_id, is_optimization_found=best_optimization is not None
             )
 
-        # If worktree mode, do not revert code and helpers,, otherwise we would have an empty diff when writing the patch in the lsp
+        # If worktree mode, do not revert code and helpers, otherwise we would have an empty diff when writing the patch in the lsp
         if self.args.worktree:
             return