Merge pull request #895 from codeflash-ai/cf-773

Run Formatter on generated tests

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
@@ -1413,11 +1413,15 @@ 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```"
+            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),