update tests

Author: KRRT7

codeflash/code_utils/formatter.py

@@ -5,7 +5,8 @@
 import subprocess
 from typing import TYPE_CHECKING
 
-import isort
+from black import Mode, format_str
+from isort import code as imports_format
 
 from codeflash.cli_cmds.console import console, logger
 
@@ -24,19 +25,31 @@ def format_code(formatter_cmds: list[str], path: Path) -> str:
     file_token = "$file"  # noqa: S105
     for command in set(formatter_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 = [
+            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)
+            result = subprocess.run(
+                formatter_cmd_list, capture_output=True, check=False
+            )
             if result.returncode == 0:
-                console.rule(f"Formatted Successfully with: {formatter_name.replace('$file', path.name)}")
+                console.rule(
+                    f"Formatted Successfully with: {formatter_name.replace('$file', path.name)}"
+                )
             else:
-                logger.error(f"Failed to format code with {' '.join(formatter_cmd_list)}")
+                logger.error(
+                    f"Failed to format code with {' '.join(formatter_cmd_list)}"
+                )
         except FileNotFoundError as e:
             from rich.panel import Panel
             from rich.text import Text
 
             panel = Panel(
-                Text.from_markup(f"⚠️  Formatter command not found: {' '.join(formatter_cmd_list)}", style="bold red"),
+                Text.from_markup(
+                    f"⚠️  Formatter command not found: {' '.join(formatter_cmd_list)}",
+                    style="bold red",
+                ),
                 expand=False,
             )
             console.print(panel)
@@ -46,12 +59,22 @@ def format_code(formatter_cmds: list[str], path: Path) -> str:
     return path.read_text(encoding="utf8")
 
 
-def sort_imports(code: str) -> str:
-    try:
-        # Deduplicate and sort imports, modify the code in memory, not on disk
-        sorted_code = isort.code(code)
-    except Exception:
-        logger.exception("Failed to sort imports with isort.")
-        return code  # Fall back to original code if isort fails
+def format_code_in_memory(code_to_format: str, *, imports_only: bool = False) -> str:
+
+    if imports_only:
+        return imports_format(code_to_format)
 
-    return sorted_code
+    formatters = [
+        (format_str, {"mode": Mode()}),
+        (imports_format, {"float_to_top": True}),
+    ]
+
+    formatted_code = code_to_format
+    try:
+        for formatter_func, formatter_kwargs in formatters:
+            formatted_code = formatter_func(formatted_code, **formatter_kwargs)
+    except (ValueError, TypeError, SyntaxError) as e:
+        logger.debug(f"Failed to format: {e}")
+        return code_to_format
+    else:
+        return formatted_code

codeflash/code_utils/instrument_existing_tests.py

@@ -4,10 +4,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
-import isort
-
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
+from codeflash.code_utils.formatter import format_code_in_memory
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import FunctionParent, TestingMode, VerificationType
 
@@ -355,8 +354,7 @@ def inject_profiling_into_existing_test(
     if test_framework == "unittest":
         new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
     tree.body = [*new_imports, create_wrapper_function(mode), *tree.body]
-    return True, isort.code(ast.unparse(tree), float_to_top=True)
-
+    return True, format_code_in_memory(ast.unparse(tree))
 
 def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.FunctionDef:
     lineno = 1

codeflash/optimization/function_optimizer.py

@@ -35,7 +35,7 @@
     N_TESTS_TO_GENERATE,
     TOTAL_LOOPING_TIME,
 )
-from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.formatter import format_code, format_code_in_memory
 from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
 from codeflash.code_utils.line_profile_utils import add_decorator_imports
 from codeflash.code_utils.remove_generated_tests import remove_functions_from_generated_tests
@@ -541,14 +541,14 @@ def reformat_code_and_helpers(
 
         new_code = format_code(self.args.formatter_cmds, path)
         if should_sort_imports:
-            new_code = sort_imports(new_code)
+            new_code = format_code_in_memory(new_code, imports_only=True)
 
         new_helper_code: dict[Path, str] = {}
         helper_functions_paths = {hf.file_path 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:
-                formatted_helper_code = sort_imports(formatted_helper_code)
+                formatted_helper_code = format_code_in_memory(formatted_helper_code, imports_only=True)
             new_helper_code[module_abspath] = formatted_helper_code
 
         return new_code, new_helper_code

tests/test_formatter.py

@@ -5,45 +5,41 @@
 import pytest
 
 from codeflash.code_utils.config_parser import parse_config_file
-from codeflash.code_utils.formatter import format_code, sort_imports
+from codeflash.code_utils.formatter import format_code, format_code_in_memory
 
 
 def test_remove_duplicate_imports():
     """Test that duplicate imports are removed when should_sort_imports is True."""
     original_code = "import os\nimport os\n"
-    new_code = sort_imports(original_code)
+    new_code = format_code_in_memory(original_code)
     assert new_code == "import os\n"
 
-
 def test_remove_multiple_duplicate_imports():
     """Test that multiple duplicate imports are removed when should_sort_imports is True."""
     original_code = "import sys\nimport os\nimport sys\n"
 
-    new_code = sort_imports(original_code)
+    new_code = format_code_in_memory(original_code)
     assert new_code == "import os\nimport sys\n"
 
-
 def test_sorting_imports():
     """Test that imports are sorted when should_sort_imports is True."""
     original_code = "import sys\nimport unittest\nimport os\n"
 
-    new_code = sort_imports(original_code)
+    new_code = format_code_in_memory(original_code)
     assert new_code == "import os\nimport sys\nimport unittest\n"
 
-
 def test_sort_imports_without_formatting():
     """Test that imports are sorted when formatting is disabled and should_sort_imports is True."""
     with tempfile.NamedTemporaryFile() as tmp:
-        tmp.write(b"import sys\nimport unittest\nimport os\n")
+        tmp.write(b"import sys\nimport unittest\nimport os\nimport sys\nimport unittest\n")
         tmp.flush()
         tmp_path = Path(tmp.name)
 
         new_code = format_code(formatter_cmds=["disabled"], path=tmp_path)
         assert new_code is not None
-        new_code = sort_imports(new_code)
+        new_code = format_code_in_memory(new_code)
         assert new_code == "import os\nimport sys\nimport unittest\n"
 
-
 def test_dedup_and_sort_imports_deduplicates():
     original_code = """
 import os
@@ -54,16 +50,15 @@ def foo():
     return os.path.join(sys.path[0], 'bar')
 """
 
-    expected = """
-import os
+    expected = """import os
 import sys
 
 
 def foo():
-    return os.path.join(sys.path[0], 'bar')
+    return os.path.join(sys.path[0], "bar")
 """
 
-    actual = sort_imports(original_code)
+    actual = format_code_in_memory(original_code)
 
     assert actual == expected
 
@@ -80,17 +75,16 @@ def foo():
     return os.path.join(sys.path[0], 'bar')
 """
 
-    expected = """
-import json
+    expected = """import json
 import os
 import sys
 
 
 def foo():
-    return os.path.join(sys.path[0], 'bar')
+    return os.path.join(sys.path[0], "bar")
 """
 
-    actual = sort_imports(original_code)
+    actual = format_code_in_memory(original_code)
 
     assert actual == expected
 

tests/test_instrument_all_and_run.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
+from codeflash.code_utils.formatter import format_code_in_memory
 from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize
 from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestType
@@ -119,10 +120,10 @@ def test_sort():
         os.chdir(original_cwd)
         assert success
         assert new_test is not None
-        assert new_test.replace('"', "'") == expected.format(
+        assert format_code_in_memory(new_test) == format_code_in_memory(expected.format(
             module_path="code_to_optimize.tests.pytest.test_perfinjector_bubble_sort_results_temp",
             tmp_dir_path=get_run_tmp_file(Path("test_return_values")),
-        ).replace('"', "'")
+        ))
 
         with test_path.open("w") as f:
             f.write(new_test)
@@ -307,9 +308,9 @@ def test_sort():
             Path(f.name), [CodePosition(7, 13), CodePosition(12, 13)], fto, Path(f.name).parent, "pytest"
         )
     assert success
-    assert new_test.replace('"', "'") == expected.format(
+    assert format_code_in_memory(new_test) == format_code_in_memory(expected.format(
         module_path=Path(f.name).name, tmp_dir_path=get_run_tmp_file(Path("test_return_values"))
-    ).replace('"', "'")
+    ))
     tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve()
     test_path = tests_root / "test_class_method_behavior_results_temp.py"
     test_path_perf = tests_root / "test_class_method_behavior_results_perf_temp.py"

tests/test_instrument_tests.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 
 from codeflash.code_utils.code_utils import get_run_tmp_file
+from codeflash.code_utils.formatter import format_code_in_memory
 from codeflash.code_utils.instrument_existing_tests import (
     FunctionImportedAsVisitor,
     inject_profiling_into_existing_test,
@@ -127,10 +128,8 @@ def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, functi
     invocation_id = f'{{line_id}}_{{codeflash_test_index}}'
     """
     if sys.version_info < (3, 12):
-        print("sys.version_info < (3, 12)")
         expected += """print(f"!######{{test_module_name}}:{{(test_class_name + '.' if test_class_name else '')}}{{test_name}}:{{function_name}}:{{loop_index}}:{{invocation_id}}######!")"""
     else:
-        print("sys.version_info >= (3, 12)")
         expected += """print(f'!######{{test_module_name}}:{{(test_class_name + '.' if test_class_name else '')}}{{test_name}}:{{function_name}}:{{loop_index}}:{{invocation_id}}######!')"""
     expected += """
     exception = None
@@ -315,6 +314,8 @@ def test_sort():
 """
         + codeflash_wrap_string
         + """
+
+
 def test_sort():
     codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
     codeflash_iteration = os.environ['CODEFLASH_TEST_ITERATION']
@@ -2845,10 +2846,8 @@ def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, functi
     codeflash_test_index = codeflash_wrap.index[test_id]
     invocation_id = f'{{line_id}}_{{codeflash_test_index}}'
 """
-    if sys.version_info < (3, 12):
-        expected += """    print(f"!######{{test_module_name}}:{{(test_class_name + '.' if test_class_name else '')}}{{test_name}}:{{function_name}}:{{loop_index}}:{{invocation_id}}######!")"""
-    else:
-        expected += """    print(f'!######{{test_module_name}}:{{(test_class_name + '.' if test_class_name else '')}}{{test_name}}:{{function_name}}:{{loop_index}}:{{invocation_id}}######!')"""
+    
+    expected += """    print(f'!######{{test_module_name}}:{{(test_class_name + '.' if test_class_name else '')}}{{test_name}}:{{function_name}}:{{loop_index}}:{{invocation_id}}######!')"""
     expected += """
     exception = None
     gc.disable()