use user pre-defined formatting commands, instead of using black

Author: mohammedahmed18

code_to_optimize/no_formatting_errors.py

@@ -1,25 +1,8 @@
-import os, sys, json, datetime, math, random
-import requests
-from collections import defaultdict, OrderedDict
-from typing import List, Dict, Optional, Union, Tuple, Any
-import numpy as np
-import pandas as pd
+import math
 
-# This is a poorly formatted Python file with many style violations
 
-
-class UnformattedExampleClass(object):
-    def __init__(
-        self,
-        name,
-        age=None,
-        email=None,
-        phone=None,
-        address=None,
-        city=None,
-        state=None,
-        zip_code=None,
-    ):
+class UnformattedExampleClass:
+    def __init__(self, name, age=None, email=None, phone=None, address=None, city=None, state=None, zip_code=None):
         self.name = name
         self.age = age
         self.email = email
@@ -40,9 +23,7 @@ def update_data(self, **kwargs):
         self.data.update(kwargs)
 
 
-def process_data(
-    data_list, filter_func=None, transform_func=None, sort_key=None, reverse=False
-):
+def process_data(data_list, filter_func=None, transform_func=None, sort_key=None, reverse=False):
     if not data_list:
         return []
     if filter_func:

codeflash/code_utils/formatter.py

@@ -1,69 +1,73 @@
 from __future__ import annotations
 
+import difflib
 import os
+import re
 import shlex
+import shutil
 import subprocess
-from typing import TYPE_CHECKING, Optional
+import tempfile
+from pathlib import Path
+from typing import Optional
 
 import isort
 
 from codeflash.cli_cmds.console import console, logger
 
-if TYPE_CHECKING:
-    from pathlib import Path
 
+def generate_unified_diff(original: str, modified: str, from_file: str, to_file: str) -> str:
+    line_pattern = re.compile(r"(.*?(?:\r\n|\n|\r|$))")
 
-def get_diff_output_by_black(filepath: str, unformatted_content: str) -> Optional[str]:
-    try:
-        from black import Mode, format_file_contents, output, report
+    def split_lines(text: str) -> list[str]:
+        lines = [match[0] for match in line_pattern.finditer(text)]
+        if lines and lines[-1] == "":
+            lines.pop()
+        return lines
 
-        formatted_content = format_file_contents(src_contents=unformatted_content, fast=True, mode=Mode())
-        return output.diff(unformatted_content, formatted_content, a_name=filepath, b_name=filepath)
-    except (ImportError, report.NothingChanged):
-        return None
+    original_lines = split_lines(original)
+    modified_lines = split_lines(modified)
 
+    diff_output = []
+    for line in difflib.unified_diff(original_lines, modified_lines, fromfile=from_file, tofile=to_file, n=5):
+        if line.endswith("\n"):
+            diff_output.append(line)
+        else:
+            diff_output.append(line + "\n")
+            diff_output.append("\\ No newline at end of file\n")
 
-def get_diff_lines_count(diff_output: str) -> int:
-    lines = diff_output.split("\n")
-
-    def is_diff_line(line: str) -> bool:
-        return line.startswith(("+", "-")) and not line.startswith(("+++", "---"))
-
-    diff_lines = [line for line in lines if is_diff_line(line)]
-    return len(diff_lines)
+    return "".join(diff_output)
 
 
-def is_safe_to_format(filepath: str, content: str, max_diff_lines: int = 100) -> bool:
-    diff_changes_str = None
-
-    diff_changes_str = get_diff_output_by_black(filepath, unformatted_content=content)
-
-    if diff_changes_str is None:
-        logger.warning("Looks like black formatter not found, make sure it is installed.")
-        return False
-
-    diff_lines_count = get_diff_lines_count(diff_changes_str)
-    if diff_lines_count > max_diff_lines:
-        logger.debug(f"Skipping formatting {filepath}: {diff_lines_count} lines would change (max: {max_diff_lines})")
-        return False
+def apply_formatter_cmds(
+    cmds: list[str],
+    path: Path,
+    test_dir_str: Optional[str],
+    print_status: bool,  # noqa
+) -> tuple[Path, str]:
+    # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
+    formatter_name = cmds[0].lower()
+    should_make_copy = False
+    file_path = path
 
-    return True
+    if test_dir_str:
+        should_make_copy = True
+        file_path = Path(test_dir_str) / "temp.py"
 
+    if not cmds or formatter_name == "disabled":
+        return path, path.read_text(encoding="utf8")
 
-def format_code(formatter_cmds: list[str], path: Path, print_status: bool = True) -> str:  # noqa
-    # TODO: Only allow a particular whitelist of formatters here to prevent arbitrary code execution
-    formatter_name = formatter_cmds[0].lower()
     if not path.exists():
-        msg = f"File {path} does not exist. Cannot format the file."
+        msg = f"File {path} does not exist. Cannot apply formatter commands."
         raise FileNotFoundError(msg)
-    file_content = path.read_text(encoding="utf8")
-    if formatter_name == "disabled" or not is_safe_to_format(filepath=str(path), content=file_content):
-        return file_content
+
+    if should_make_copy:
+        shutil.copy2(path, file_path)
 
     file_token = "$file"  # noqa: S105
-    for command in formatter_cmds:
+
+    for command in cmds:
         formatter_cmd_list = shlex.split(command, posix=os.name != "nt")
-        formatter_cmd_list = [path.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
+        formatter_cmd_list = [file_path.as_posix() if chunk == file_token else chunk for chunk in formatter_cmd_list]
         try:
             result = subprocess.run(formatter_cmd_list, capture_output=True, check=False)
             if result.returncode == 0:
@@ -83,7 +87,45 @@ def format_code(formatter_cmds: list[str], path: Path, print_status: bool = True
 
             raise e from None
 
-    return path.read_text(encoding="utf8")
+    return file_path, file_path.read_text(encoding="utf8")
+
+
+def get_diff_lines_count(diff_output: str) -> int:
+    lines = diff_output.split("\n")
+
+    def is_diff_line(line: str) -> bool:
+        return line.startswith(("+", "-")) and not line.startswith(("+++", "---"))
+
+    diff_lines = [line for line in lines if is_diff_line(line)]
+    return len(diff_lines)
+
+
+def format_code(formatter_cmds: list[str], path: Path, optimized_function: str = "", print_status: bool = True) -> str:  # noqa
+    with tempfile.TemporaryDirectory() as test_dir_str:
+        max_diff_lines = 100
+
+        original_code = path.read_text(encoding="utf8")
+        # we dont' count the formatting diff for the optimized function as it should be well-formatted (if it's provided)
+        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
+
+        _, 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}")
+        return formatted_code
 
 
 def sort_imports(code: str) -> str:

codeflash/optimization/function_optimizer.py

@@ -302,7 +302,10 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
                 )
 
                 new_code, new_helper_code = self.reformat_code_and_helpers(
-                    code_context.helper_functions, explanation.file_path, self.function_to_optimize_source_code
+                    code_context.helper_functions,
+                    explanation.file_path,
+                    self.function_to_optimize_source_code,
+                    optimized_function=best_optimization.candidate.source_code,
                 )
 
                 existing_tests = existing_tests_source_for(
@@ -591,18 +594,18 @@ def write_code_and_helpers(original_code: str, original_helper_code: dict[Path,
                 f.write(helper_code)
 
     def reformat_code_and_helpers(
-        self, helper_functions: list[FunctionSource], path: Path, original_code: str
+        self, helper_functions: list[FunctionSource], path: Path, original_code: str, optimized_function: str
     ) -> tuple[str, dict[Path, str]]:
         should_sort_imports = not self.args.disable_imports_sorting
         if should_sort_imports and isort.code(original_code) != original_code:
             should_sort_imports = False
 
-        new_code = format_code(self.args.formatter_cmds, path)
+        new_code = format_code(self.args.formatter_cmds, path, optimized_function=optimized_function)
         if should_sort_imports:
             new_code = sort_imports(new_code)
 
         new_helper_code: dict[Path, str] = {}
-        helper_functions_paths = {hf.file_path for hf in helper_functions}
+        helper_functions_paths = {hf.source_code for hf in helper_functions}
         for module_abspath in helper_functions_paths:
             formatted_helper_code = format_code(self.args.formatter_cmds, module_abspath)
             if should_sort_imports:

poetry.lock

@@ -93,7 +93,6 @@ crosshair-tool = ">=0.0.78"
 coverage = ">=7.6.4"
 line_profiler=">=4.2.0" #this is the minimum version which supports python 3.13
 platformdirs = ">=4.3.7"
-black = "^25.1.0"
 [tool.poetry.group.dev]
 optional = true
 

pyproject.toml

@@ -263,6 +263,16 @@ def _run_formatting_test(source_filename: str, should_content_change: bool):
             helper_functions=[],
             path=target_path,
             original_code=optimizer.function_to_optimize_source_code,
+            # this is just for testing, but in practice, this would be an optimized function code and it will be well-formatted
+            optimized_function="""    def process(self):
+        data=self.load_data()
+        if not data:return{"success":False,"error":"No data loaded"}
+        
+        validated_data=self.validate_data(data)
+        processed_result=process_data(validated_data,
+                                    filter_func=lambda x:x.get('active',True),
+                                    transform_func=lambda x:{**x,'processed_at':datetime.datetime.now().isoformat()},
+                                    sort_key=lambda x:x.get('name',''))""",
         )
 
         if should_content_change: