Merge remote-tracking branch 'origin/main' into cf-835

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/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
@@ -1417,11 +1417,15 @@ def process_review(
         generated_tests_str = ""
         for test in generated_tests.generated_tests:
             if map_gen_test_file_to_no_of_tests[test.behavior_file_path] > 0:
-                generated_tests_str += f"```python\n{test.generated_original_test_source}\n```"
+                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),
@@ -1569,8 +1573,7 @@ def establish_original_code_baseline(
     ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]:
         line_profile_results = {"timings": {}, "unit": 0, "str_out": ""}
         # For the original function - run the tests and get the runtime, plus coverage
-        test_framework = self.args.test_framework
-        assert test_framework in {"pytest", "unittest"}
+        assert (test_framework := self.args.test_framework) in {"pytest", "unittest"}  # noqa: RUF018
         success = True
 
         test_env = self.get_test_env(codeflash_loop_index=0, codeflash_test_iteration=0, codeflash_tracer_disable=1)
@@ -1748,8 +1751,7 @@ def run_optimized_candidate(
         original_helper_code: dict[Path, str],
         file_path_to_helper_classes: dict[Path, set[str]],
     ) -> Result[OptimizedCandidateResult, str]:
-        test_framework = self.args.test_framework
-        assert test_framework in {"pytest", "unittest"}
+        assert (test_framework := self.args.test_framework) in {"pytest", "unittest"}  # noqa: RUF018
 
         with progress_bar("Testing optimization candidate"):
             test_env = self.get_test_env(