tests: edge cases for cst formatting

Author: mohammedahmed18

codeflash/discovery/functions_to_optimize.py

@@ -80,27 +80,29 @@ def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
                     ending_line=pos.end.line,
                 )
             )
-
 class CodeRangeFunctionVisitor(cst.CSTVisitor):
-    METADATA_DEPENDENCIES = (cst.metadata.PositionProvider, cst.metadata.QualifiedNameProvider)
-    
+    METADATA_DEPENDENCIES = (
+        cst.metadata.PositionProvider,
+        cst.metadata.QualifiedNameProvider,
+    )
+
     def __init__(self, target_function_name: str) -> None:
         super().__init__()
         self.target_func = target_function_name
-        self.current_path = []
         self.start_line: Optional[int] = None
         self.end_line: Optional[int] = None
 
     def visit_FunctionDef(self, node: cst.FunctionDef) -> bool:
-        qualified_names = {
-            str(qn.name) for qn in 
+        qualified_names = [
+            str(qn.name).replace(".<locals>", "") for qn in 
             self.get_metadata(cst.metadata.QualifiedNameProvider, node)
-        }
-        
+        ]
         if self.target_func in qualified_names:
-            position = self.get_metadata(cst.metadata.PositionProvider, node)
-            self.start_line = position.start.line
-            self.end_line = position.end.line
+            func_position = self.get_metadata(cst.metadata.PositionProvider, node)
+            decorators_count = len(node.decorators)
+            self.start_line = func_position.start.line - decorators_count
+            self.end_line = func_position.end.line
+            return False
 
 
 class FunctionWithReturnStatement(ast.NodeVisitor):

codeflash/optimization/function_optimizer.py

@@ -302,13 +302,6 @@ def optimize_function(self) -> Result[BestOptimization, str]:  # noqa: PLR0911
                     code_context=code_context, optimized_code=best_optimization.candidate.source_code
                 )
 
-                print("file_path_to_helper_classes\n", file_path_to_helper_classes)
-                filepaths_to_inspect = [
-                    self.function_to_optimize.file_path,
-                    *list({helper.file_path for helper in code_context.helper_functions}),
-                ]
-                print("filepaths_to_inspect\n", filepaths_to_inspect)
-                
                 new_code, new_helper_code = self.reformat_code_and_helpers(
                     code_context.helper_functions, explanation.file_path, self.function_to_optimize_source_code,
                     opt_func_name=explanation.function_name
@@ -610,13 +603,20 @@ def reformat_code_and_helpers(
             should_sort_imports = False
 
         whole_file_content = path.read_text(encoding="utf8")
-        wrapper = cst.metadata.MetadataWrapper(cst.parse_module(whole_file_content))
+        wrapper: cst.metadata.MetadataWrapper | None = None
+        try:
+            wrapper = cst.metadata.MetadataWrapper(cst.parse_module(whole_file_content))
+        except cst.ParserSyntaxError as e:
+            logger.error(f"Syntax error detected, aborting reformatting.")
+            return original_code, {}
+
         visitor = CodeRangeFunctionVisitor(target_function_name=opt_func_name)
         wrapper.visit(visitor)
 
         lines = whole_file_content.splitlines(keepends=True)
         if visitor.start_line == None:
             logger.error(f"Could not find function {opt_func_name} in {path}, aborting reformatting.")
+            return original_code, {}
         else:
             opt_func_source_lines = lines[visitor.start_line-1:visitor.end_line]
 

tests/test_formatter.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 import tempfile
 from pathlib import Path
@@ -7,6 +8,9 @@
 from codeflash.code_utils.config_parser import parse_config_file
 from codeflash.code_utils.formatter import format_code, sort_imports
 
+from codeflash.discovery.functions_to_optimize import FunctionToOptimize
+from codeflash.optimization.function_optimizer import FunctionOptimizer
+from codeflash.verification.verification_utils import TestConfig
 
 def test_remove_duplicate_imports():
     """Test that duplicate imports are removed when should_sort_imports is True."""
@@ -209,3 +213,255 @@ def foo():
         tmp_path = tmp.name
         with pytest.raises(FileNotFoundError):
             format_code(formatter_cmds=["exit 1"], path=Path(tmp_path))
+
+############################################################
+################ CST based formatting tests ################
+############################################################
+@pytest.fixture
+def setup_cst_formatter_args():
+    """Common setup for reformat_code_and_helpers tests."""
+    def _setup(unformatted_code, function_name):
+        test_dir = Path(tempfile.mkdtemp())
+        target_path = test_dir / "target.py"
+        target_path.write_text(unformatted_code, encoding="utf-8")
+        
+        function_to_optimize = FunctionToOptimize(
+            function_name=function_name, 
+            parents=[], 
+            file_path=target_path
+        )
+        
+        test_cfg = TestConfig(
+            tests_root=test_dir,
+            project_root_path=test_dir,
+            test_framework="pytest",
+            tests_project_rootdir=test_dir,
+        )
+        
+        args = argparse.Namespace(
+            disable_imports_sorting=False,
+            formatter_cmds=[
+                "ruff check --exit-zero --fix $file",
+                "ruff format $file"
+            ],
+        )
+        
+        optimizer = FunctionOptimizer(
+            function_to_optimize=function_to_optimize,
+            test_cfg=test_cfg,
+            args=args,
+        )
+        
+        return optimizer, target_path, function_to_optimize
+    
+    yield _setup
+
+
+def test_reformat_code_and_helpers(setup_cst_formatter_args):
+    """
+    reformat_code_and_helpers should only format the code that is optimized not the whole file, to avoid large diffing
+    """
+    unformatted_code = """import sys
+
+
+def lol():
+    print(       "lol" )
+
+
+
+
+class MyClass:
+    def __init__(self, x=0):
+        self.x = x
+
+    def lol(self):
+        print(       "lol" )
+
+    def lol2    (self):
+        print(       " lol2" )"""
+    
+    expected_code = """import sys
+
+
+def lol():
+    print(       "lol" )
+
+
+
+
+class MyClass:
+    def __init__(self, x=0):
+        self.x = x
+
+    def lol(self):
+        print(       "lol" )
+
+    def lol2(self):
+        print(" lol2")
+"""
+    
+    optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+        unformatted_code, "MyClass.lol2"
+    )
+    
+    formatted_code, _ = optimizer.reformat_code_and_helpers(
+        helper_functions=[],
+        path=target_path,
+        original_code=optimizer.function_to_optimize_source_code,
+        opt_func_name=function_to_optimize.function_name
+    )
+    
+    assert formatted_code == expected_code
+
+
+def test_reformat_code_and_helpers_with_duplicated_target_function_names(setup_cst_formatter_args):
+    unformatted_code = """import sys
+def lol():
+    print(       "lol" )
+
+class MyClass:
+    def __init__(self, x=0):
+        self.x = x
+
+    def lol(self):
+        print(       "lol" )"""
+    
+    expected_code = """import sys
+def lol():
+    print(       "lol" )
+
+class MyClass:
+    def __init__(self, x=0):
+        self.x = x
+
+    def lol(self):
+        print("lol")
+"""
+    
+    optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+        unformatted_code, "MyClass.lol"
+    )
+    
+    formatted_code, _ = optimizer.reformat_code_and_helpers(
+        helper_functions=[],
+        path=target_path,
+        original_code=optimizer.function_to_optimize_source_code,
+        opt_func_name=function_to_optimize.function_name
+    )
+    
+    assert formatted_code == expected_code
+
+
+
+def test_formatting_nested_functions(setup_cst_formatter_args):
+    unformatted_code = """def hello():
+    print("Hello")
+    def nested_function()   :
+        print       ("This is a nested function")
+    def another_nested_function():
+        print   ("This is another nested function")"""
+    
+    expected_code = """def hello():
+    print("Hello")
+    def nested_function():
+        print("This is a nested function")
+    def another_nested_function():
+        print   ("This is another nested function")"""
+    
+    optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+        unformatted_code, "hello.nested_function"
+    )
+    
+    formatted_code, _ = optimizer.reformat_code_and_helpers(
+        helper_functions=[],
+        path=target_path,
+        original_code=optimizer.function_to_optimize_source_code,
+        opt_func_name=function_to_optimize.function_name
+    )
+    
+    assert formatted_code == expected_code
+
+
+def test_formatting_standalone_functions(setup_cst_formatter_args):
+    unformatted_code = """def func1   ():
+    print(      "This is a function with bad formatting")
+def func2() :
+    print   (   "This is another function with bad formatting"  )
+"""
+    
+    expected_code = """def func1   ():
+    print(      "This is a function with bad formatting")
+def func2():
+    print("This is another function with bad formatting")
+"""
+    
+    optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+        unformatted_code, "func2"
+    )
+    
+    formatted_code, _ = optimizer.reformat_code_and_helpers(
+        helper_functions=[],
+        path=target_path,
+        original_code=optimizer.function_to_optimize_source_code,
+        opt_func_name=function_to_optimize.function_name
+    )
+    
+    assert formatted_code == expected_code
+
+
+def test_formatting_function_with_decorators(setup_cst_formatter_args):
+   unformatted_code = """@decorator1
+@decorator2(   arg1 , arg2   )
+def func1   ():
+   print(      "This is a function with bad formatting")
+
+@another_decorator(     arg)
+def func2   (  x,y  ):
+   print   (   "This is another function with bad formatting"  )"""
+   
+   expected_code = """@decorator1
+@decorator2(   arg1 , arg2   )
+def func1   ():
+   print(      "This is a function with bad formatting")
+
+@another_decorator(arg)
+def func2(x, y):
+    print("This is another function with bad formatting")
+"""
+   
+   optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+       unformatted_code, "func2"
+   )
+   
+   formatted_code, _ = optimizer.reformat_code_and_helpers(
+       helper_functions=[],
+       path=target_path,
+       original_code=optimizer.function_to_optimize_source_code,
+       opt_func_name=function_to_optimize.function_name
+   )
+   
+   assert formatted_code == expected_code
+
+
+def test_formatting_function_with_syntax_error(setup_cst_formatter_args):
+    """shouldn't happen anyway, but just in case"""
+    unformatted_code = """def func1():
+    print("This is a function with a syntax error"
+def func2():
+    print("This is another function with a syntax error")
+"""
+    
+    expected_code = unformatted_code  # No formatting should be applied due to syntax error
+    
+    optimizer, target_path, function_to_optimize = setup_cst_formatter_args(
+        unformatted_code, "func2"
+    )
+    
+    formatted_code, _ = optimizer.reformat_code_and_helpers(
+        helper_functions=[],
+        path=target_path,
+        original_code=optimizer.function_to_optimize_source_code,
+        opt_func_name=function_to_optimize.function_name
+    )
+    
+    assert formatted_code == expected_code