wip

Author: Codeflash Bot

codeflash/code_utils/formatter.py

@@ -96,6 +96,20 @@ def is_diff_line(line: str) -> bool:
     return len(diff_lines)
 
 
+def format_generated_code(generated_test_source: str, formatter_cmds: Union[list[str], None] = None) -> str:
+    formatter_name = formatter_cmds[0].lower() if formatter_cmds else "disabled"
+    if formatter_name == "disabled":
+        return re.sub(r"\n{2,}", "\n\n", generated_test_source)
+    # try running formatter, if nothing changes (could be due to formatting failing or no actual formatting needed)
+    original_temp, test_dir_str, exit_on_failure = None, None, True
+    formatted_temp, formatted_code, changed = apply_formatter_cmds(
+        formatter_cmds, original_temp, test_dir_str, print_status=False, exit_on_failure=exit_on_failure
+    )
+    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 +134,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"