enhancements

Author: mohammedahmed18

codeflash/code_utils/formatter.py

@@ -104,32 +104,39 @@ def format_code(
     formatter_cmds: list[str],
     path: Union[str, Path],
     optimized_function: str = "",
+    check_diff: bool = False,  # noqa
     print_status: bool = True,  # noqa
 ) -> str:
     with tempfile.TemporaryDirectory() as test_dir_str:
-        max_diff_lines = 100
-
         if isinstance(path, str):
             path = Path(path)
 
         original_code = path.read_text(encoding="utf8")
-        # we dont' count the formatting diff for the optimized function as it should be well-formatted
-        original_code_without_opfunc = original_code.replace(optimized_function, "")
-
-        original_temp = Path(test_dir_str) / "original_temp.py"
-        original_temp.write_text(original_code_without_opfunc, encoding="utf8")
-
-        formatted_temp, formatted_code = apply_formatter_cmds(
-            formatter_cmds, original_temp, test_dir_str, print_status=False
-        )
-
-        diff_output = generate_unified_diff(
-            original_code_without_opfunc, formatted_code, from_file=str(original_temp), to_file=str(formatted_temp)
-        )
-        diff_lines_count = get_diff_lines_count(diff_output)
-        if diff_lines_count > max_diff_lines:
-            logger.debug(f"Skipping formatting {path}: {diff_lines_count} lines would change (max: {max_diff_lines})")
-            return original_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
+            original_code_without_opfunc = original_code.replace(optimized_function, "")
+
+            original_temp = Path(test_dir_str) / "original_temp.py"
+            original_temp.write_text(original_code_without_opfunc, encoding="utf8")
+
+            formatted_temp, formatted_code = apply_formatter_cmds(
+                formatter_cmds, original_temp, test_dir_str, print_status=False
+            )
+
+            diff_output = generate_unified_diff(
+                original_code_without_opfunc, formatted_code, from_file=str(original_temp), to_file=str(formatted_temp)
+            )
+            diff_lines_count = get_diff_lines_count(diff_output)
+
+            max_diff_lines = min(int(original_code_lines * 0.3), 50)
+
+            if diff_lines_count > max_diff_lines and max_diff_lines != -1:
+                logger.debug(
+                    f"Skipping formatting {path}: {diff_lines_count} lines would change (max: {max_diff_lines})"
+                )
+                return original_code
 
         _, formatted_code = apply_formatter_cmds(formatter_cmds, path, test_dir_str=None, print_status=print_status)
         logger.debug(f"Formatted {path} with commands: {formatter_cmds}")

codeflash/optimization/function_optimizer.py

@@ -612,7 +612,7 @@ def reformat_code_and_helpers(
         if should_sort_imports and isort.code(original_code) != original_code:
             should_sort_imports = False
 
-        new_code = format_code(self.args.formatter_cmds, path, optimized_function=optimized_function)
+        new_code = format_code(self.args.formatter_cmds, path, optimized_function=optimized_function, check_diff=True)
         if should_sort_imports:
             new_code = sort_imports(new_code)
 
@@ -621,7 +621,7 @@ def reformat_code_and_helpers(
             module_abspath = hp.file_path
             hp_source_code = hp.source_code
             formatted_helper_code = format_code(
-                self.args.formatter_cmds, module_abspath, optimized_function=hp_source_code
+                self.args.formatter_cmds, module_abspath, optimized_function=hp_source_code, check_diff=True
             )
             if should_sort_imports:
                 formatted_helper_code = sort_imports(formatted_helper_code)