refactoring

Author: mohammedahmed18

code_to_optimize/few_formatting_errors.py

@@ -1,6 +1,6 @@
 import os
 
-class BadlyFormattedClass(object):
+class UnformattedExampleClass(object):
     def __init__(
         self,
         name,

code_to_optimize/many_formatting_errors.py

@@ -3,7 +3,7 @@
 
 # This is a poorly formatted Python file with many style violations
 
-class   BadlyFormattedClass( object ):
+class   UnformattedExampleClass( object ):
     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;self.phone=phone
         self.address=address;self.city=city;self.state=state;self.zip_code=zip_code
@@ -84,7 +84,7 @@ def generate_report(data,include_stats=True,include_charts=False,format_type='js
 
     return result
 
-class   DataProcessor  (  BadlyFormattedClass  )  :
+class   DataProcessor  (  UnformattedExampleClass  )  :
     def __init__(self,data_source,config=None,debug=False):
         super().__init__("DataProcessor")
         self.data_source=data_source;self.config=config or{};self.debug=debug

codeflash/code_utils/formatter.py

@@ -3,7 +3,7 @@
 import os
 import shlex
 import subprocess
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import isort
 
@@ -12,45 +12,68 @@
 if TYPE_CHECKING:
     from pathlib import Path
 
-
-def should_format_file(filepath, max_lines_changed=100):
-        try:
-            # check if black is installed
-            subprocess.run(['black', '--version'], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-
-            result = subprocess.run(
-                ['black', '--diff', filepath], 
-                capture_output=True, 
-                text=True
-            )
-                
-            diff_lines = [line for line in result.stdout.split('\n') 
-                        if line.startswith(('+', '-')) and not line.startswith(('+++', '---'))]
-            
-            changes_count = len(diff_lines)
-            
-            if changes_count > max_lines_changed:
-                logger.debug(f"Skipping {filepath}: {changes_count} lines would change (max: {max_lines_changed})")
-                return False
-            
-            return True
-            
-        except subprocess.CalledProcessError:
-            logger.warning(f"black --diff command failed for {filepath}")
-            return False
-        except FileNotFoundError:
-            logger.warning("black formatter is not installed. Skipping formatting diff check.")
-            return False
-
-
+def get_diff_lines_output_by_black(filepath: str) -> Optional[str]:
+    try:
+        subprocess.run(['black', '--version'], check=True,
+                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        result = subprocess.run(
+            ['black', '--diff', filepath],
+            capture_output=True,
+            text=True
+        )
+        return result.stdout.strip() if result.stdout else None
+    except (FileNotFoundError):
+        return None
+
+
+def get_diff_lines_output_by_ruff(filepath: str) -> Optional[str]:
+    try:
+        subprocess.run(['ruff', '--version'], check=True,
+                     stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        result = subprocess.run(
+            ['ruff', "format", '--diff', filepath],
+            capture_output=True,
+            text=True
+        )
+        return result.stdout.strip() if result.stdout else None
+    except (FileNotFoundError):
+        return None
+
+
+def get_diff_lines_count(diff_output: str) -> int:
+    diff_lines = [line for line in diff_output.split('\n') 
+                  if line.startswith(('+', '-')) and not line.startswith(('+++', '---'))]
+    return len(diff_lines)
+
+def is_safe_to_format(filepath: str, max_diff_lines: int = 100) -> bool:
+    diff_changes_stdout = None
+
+    diff_changes_stdout = get_diff_lines_output_by_black(filepath)
+
+    if diff_changes_stdout is None:
+        logger.warning(f"black formatter not found, trying ruff instead...")
+        diff_changes_stdout = get_diff_lines_output_by_ruff(filepath)
+        if diff_changes_stdout is None:
+            msg = f"Both ruff, black formatters not found, skipping formatting diff check."
+            logger.warning(msg)
+            raise FileNotFoundError(msg)
+    
+    diff_lines_count = get_diff_lines_count(diff_changes_stdout)
+    
+    if diff_lines_count > max_diff_lines:
+        logger.debug(f"Skipping {filepath}: {diff_lines_count} lines would change (max: {max_diff_lines})")
+        return False
+    else:
+        return True
+        
 
 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."
         raise FileNotFoundError(msg)
-    if formatter_name == "disabled" or not should_format_file(path):
+    if formatter_name == "disabled" or not is_safe_to_format(path):     # few -> False, large -> True
         return path.read_text(encoding="utf8")
 
     file_token = "$file"  # noqa: S105

tests/test_formatter.py

@@ -268,12 +268,19 @@ def _run_formatting_test(source_filename: str, should_content_change: bool):
         else:
             assert content == original, f"Expected content to remain unchanged for {source_filename}"
 
+def _ruff_or_black_installed() -> bool:
+    return shutil.which("black") is not None or shutil.which("ruff") is not None
+
 
 def test_formatting_file_with_many_diffs():
     """Test that files with many formatting errors are skipped (content unchanged)."""
+    if not _ruff_or_black_installed():
+        pytest.skip("Neither black nor ruff is installed, skipping formatting tests.")
     _run_formatting_test("many_formatting_errors.py", should_content_change=False)
 
 
 def test_formatting_file_with_few_diffs():
     """Test that files with few formatting errors are formatted (content changed)."""
+    if not _ruff_or_black_installed():
+        pytest.skip("Neither black nor ruff is installed, skipping formatting tests.")
     _run_formatting_test("few_formatting_errors.py", should_content_change=True)