From 325ec7d7417bb8af3e53fe392ac072b5b555d2d1 Mon Sep 17 00:00:00 2001 From: KRRT7 Date: Wed, 18 Feb 2026 06:30:16 +0000 Subject: [PATCH 1/5] refactor: inline async decorators to remove codeflash import dependency Instead of injecting `from codeflash.code_utils.codeflash_wrap_decorator import ...` into instrumented source files, inject the decorator function definitions directly. This removes the hard dependency on the codeflash package being importable at runtime in the target environment, matching the pattern already used for sync instrumentation. --- .../code_utils/instrument_existing_tests.py | 235 ++++++++++++++++-- tests/test_async_run_and_parse_tests.py | 84 ++----- tests/test_instrument_async_tests.py | 220 ++++++++-------- 3 files changed, 363 insertions(+), 176 deletions(-) diff --git a/codeflash/code_utils/instrument_existing_tests.py b/codeflash/code_utils/instrument_existing_tests.py index 4366468d0..8c53c8e01 100644 --- a/codeflash/code_utils/instrument_existing_tests.py +++ b/codeflash/code_utils/instrument_existing_tests.py @@ -1497,15 +1497,218 @@ def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Ca return False -class AsyncDecoratorImportAdder(cst.CSTTransformer): - """Transformer that adds the import for async decorators.""" +def get_behavior_async_inline_code() -> str: + return """import asyncio +import gc +import os +import sqlite3 +from functools import wraps +from pathlib import Path +from tempfile import TemporaryDirectory + +import dill as pickle + + +def get_run_tmp_file(file_path): + if not hasattr(get_run_tmp_file, "tmpdir"): + get_run_tmp_file.tmpdir = TemporaryDirectory(prefix="codeflash_") + return Path(get_run_tmp_file.tmpdir.name) / file_path + + +def extract_test_context_from_env(): + test_module = os.environ["CODEFLASH_TEST_MODULE"] + test_class = os.environ.get("CODEFLASH_TEST_CLASS", None) + test_function = os.environ["CODEFLASH_TEST_FUNCTION"] + if test_module and test_function: + return (test_module, test_class if test_class else None, test_function) + raise RuntimeError( + "Test context environment variables not set - ensure tests are run through codeflash test runner" + ) + + +def codeflash_behavior_async(func): + @wraps(func) + async def async_wrapper(*args, **kwargs): + loop = asyncio.get_running_loop() + function_name = func.__name__ + line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"] + loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"]) + test_module_name, test_class_name, test_name = extract_test_context_from_env() + test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}" + if not hasattr(async_wrapper, "index"): + async_wrapper.index = {} + if test_id in async_wrapper.index: + async_wrapper.index[test_id] += 1 + else: + async_wrapper.index[test_id] = 0 + codeflash_test_index = async_wrapper.index[test_id] + invocation_id = f"{line_id}_{codeflash_test_index}" + class_prefix = (test_class_name + ".") if test_class_name else "" + test_stdout_tag = f"{test_module_name}:{class_prefix}{test_name}:{function_name}:{loop_index}:{invocation_id}" + print(f"!$######{test_stdout_tag}######$!") + iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0") + db_path = get_run_tmp_file(Path(f"test_return_values_{iteration}.sqlite")) + codeflash_con = sqlite3.connect(db_path) + codeflash_cur = codeflash_con.cursor() + codeflash_cur.execute( + "CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, " + "test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, " + "runtime INTEGER, return_value BLOB, verification_type TEXT)" + ) + exception = None + counter = loop.time() + gc.disable() + try: + ret = func(*args, **kwargs) + counter = loop.time() + return_value = await ret + codeflash_duration = int((loop.time() - counter) * 1_000_000_000) + except Exception as e: + codeflash_duration = int((loop.time() - counter) * 1_000_000_000) + exception = e + finally: + gc.enable() + print(f"!######{test_stdout_tag}######!") + pickled_return_value = pickle.dumps(exception) if exception else pickle.dumps((args, kwargs, return_value)) + codeflash_cur.execute( + "INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ( + test_module_name, + test_class_name, + test_name, + function_name, + loop_index, + invocation_id, + codeflash_duration, + pickled_return_value, + "function_call", + ), + ) + codeflash_con.commit() + codeflash_con.close() + if exception: + raise exception + return return_value + return async_wrapper +""" + + +def get_performance_async_inline_code() -> str: + return """import asyncio +import gc +import os +from functools import wraps + + +def extract_test_context_from_env(): + test_module = os.environ["CODEFLASH_TEST_MODULE"] + test_class = os.environ.get("CODEFLASH_TEST_CLASS", None) + test_function = os.environ["CODEFLASH_TEST_FUNCTION"] + if test_module and test_function: + return (test_module, test_class if test_class else None, test_function) + raise RuntimeError( + "Test context environment variables not set - ensure tests are run through codeflash test runner" + ) + + +def codeflash_performance_async(func): + @wraps(func) + async def async_wrapper(*args, **kwargs): + loop = asyncio.get_running_loop() + function_name = func.__name__ + line_id = os.environ["CODEFLASH_CURRENT_LINE_ID"] + loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"]) + test_module_name, test_class_name, test_name = extract_test_context_from_env() + test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}" + if not hasattr(async_wrapper, "index"): + async_wrapper.index = {} + if test_id in async_wrapper.index: + async_wrapper.index[test_id] += 1 + else: + async_wrapper.index[test_id] = 0 + codeflash_test_index = async_wrapper.index[test_id] + invocation_id = f"{line_id}_{codeflash_test_index}" + class_prefix = (test_class_name + ".") if test_class_name else "" + test_stdout_tag = f"{test_module_name}:{class_prefix}{test_name}:{function_name}:{loop_index}:{invocation_id}" + print(f"!$######{test_stdout_tag}######$!") + exception = None + counter = loop.time() + gc.disable() + try: + ret = func(*args, **kwargs) + counter = loop.time() + return_value = await ret + codeflash_duration = int((loop.time() - counter) * 1_000_000_000) + except Exception as e: + codeflash_duration = int((loop.time() - counter) * 1_000_000_000) + exception = e + finally: + gc.enable() + print(f"!######{test_stdout_tag}:{codeflash_duration}######!") + if exception: + raise exception + return return_value + return async_wrapper +""" + + +def get_concurrency_async_inline_code() -> str: + return """import asyncio +import gc +import os +import time +from functools import wraps + + +def codeflash_concurrency_async(func): + @wraps(func) + async def async_wrapper(*args, **kwargs): + function_name = func.__name__ + concurrency_factor = int(os.environ.get("CODEFLASH_CONCURRENCY_FACTOR", "10")) + test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "") + test_class_name = os.environ.get("CODEFLASH_TEST_CLASS", "") + test_function = os.environ.get("CODEFLASH_TEST_FUNCTION", "") + loop_index = os.environ.get("CODEFLASH_LOOP_INDEX", "0") + gc.disable() + try: + seq_start = time.perf_counter_ns() + for _ in range(concurrency_factor): + result = await func(*args, **kwargs) + sequential_time = time.perf_counter_ns() - seq_start + finally: + gc.enable() + gc.disable() + try: + conc_start = time.perf_counter_ns() + tasks = [func(*args, **kwargs) for _ in range(concurrency_factor)] + await asyncio.gather(*tasks) + concurrent_time = time.perf_counter_ns() - conc_start + finally: + gc.enable() + tag = f"{test_module_name}:{test_class_name}:{test_function}:{function_name}:{loop_index}" + print(f"!@######CONC:{tag}:{sequential_time}:{concurrent_time}:{concurrency_factor}######@!") + return result + return async_wrapper +""" + + +def get_async_inline_code(mode: TestingMode) -> str: + if mode == TestingMode.BEHAVIOR: + return get_behavior_async_inline_code() + if mode == TestingMode.CONCURRENCY: + return get_concurrency_async_inline_code() + return get_performance_async_inline_code() + + +class AsyncInlineCodeInjector(cst.CSTTransformer): + """Injects async decorator function definitions inline instead of importing from codeflash.""" def __init__(self, mode: TestingMode = TestingMode.BEHAVIOR) -> None: self.mode = mode - self.has_import = False + self.has_inline_definition = False + self.has_old_import = False def _get_decorator_name(self) -> str: - """Get the decorator name based on the testing mode.""" if self.mode == TestingMode.BEHAVIOR: return "codeflash_behavior_async" if self.mode == TestingMode.CONCURRENCY: @@ -1513,7 +1716,6 @@ def _get_decorator_name(self) -> str: return "codeflash_performance_async" def visit_ImportFrom(self, node: cst.ImportFrom) -> None: - # Check if the async decorator import is already present if ( isinstance(node.module, cst.Attribute) and isinstance(node.module.value, cst.Attribute) @@ -1526,21 +1728,18 @@ def visit_ImportFrom(self, node: cst.ImportFrom) -> None: decorator_name = self._get_decorator_name() for import_alias in node.names: if import_alias.name.value == decorator_name: - self.has_import = True + self.has_old_import = True + + def visit_FunctionDef(self, node: cst.FunctionDef) -> None: + if node.name.value == self._get_decorator_name(): + self.has_inline_definition = True def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: - # If the import is already there, don't add it again - if self.has_import: + if self.has_inline_definition or self.has_old_import: return updated_node - - # Choose import based on mode - decorator_name = self._get_decorator_name() - - # Parse the import statement into a CST node - import_node = cst.parse_statement(f"from codeflash.code_utils.codeflash_wrap_decorator import {decorator_name}") - - # Add the import to the module's body - return updated_node.with_changes(body=[import_node, *list(updated_node.body)]) + inline_code = get_async_inline_code(self.mode) + inline_stmts = cst.parse_module(inline_code).body + return updated_node.with_changes(body=[*inline_stmts, *list(updated_node.body)]) def add_async_decorator_to_function( @@ -1575,7 +1774,7 @@ def add_async_decorator_to_function( # Add the import if decorator was added if decorator_transformer.added_decorator: - import_transformer = AsyncDecoratorImportAdder(mode) + import_transformer = AsyncInlineCodeInjector(mode) module = module.visit(import_transformer) modified_code = sort_imports(code=module.code, float_to_top=True) diff --git a/tests/test_async_run_and_parse_tests.py b/tests/test_async_run_and_parse_tests.py index 1eb667b3f..5750a015f 100644 --- a/tests/test_async_run_and_parse_tests.py +++ b/tests/test_async_run_and_parse_tests.py @@ -9,6 +9,7 @@ from codeflash.code_utils.instrument_existing_tests import ( add_async_decorator_to_function, + get_async_inline_code, inject_profiling_into_existing_test, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -59,12 +60,16 @@ async def test_async_sort(): assert source_success - # Verify the file was modified + # Verify the file was modified with exact expected output instrumented_source = fto_path.read_text("utf-8") - assert ( - '''import asyncio\nfrom typing import List, Union\n\nfrom codeflash.code_utils.codeflash_wrap_decorator import \\\n codeflash_behavior_async\n\n\n@codeflash_behavior_async\nasync def async_sorter(lst: List[Union[int, float]]) -> List[Union[int, float]]:\n """\n Async bubble sort implementation for testing.\n """\n print("codeflash stdout: Async sorting list")\n \n await asyncio.sleep(0.01)\n \n n = len(lst)\n for i in range(n):\n for j in range(0, n - i - 1):\n if lst[j] > lst[j + 1]:\n lst[j], lst[j + 1] = lst[j + 1], lst[j]\n \n result = lst.copy()\n print(f"result: {result}")\n return result\n\n\nclass AsyncBubbleSorter:\n """Class with async sorting method for testing."""\n \n async def sorter(self, lst: List[Union[int, float]]) -> List[Union[int, float]]:\n """\n Async bubble sort implementation within a class.\n """\n print("codeflash stdout: AsyncBubbleSorter.sorter() called")\n \n # Add some async delay\n await asyncio.sleep(0.005)\n \n n = len(lst)\n for i in range(n):\n for j in range(0, n - i - 1):\n if lst[j] > lst[j + 1]:\n lst[j], lst[j + 1] = lst[j + 1], lst[j]\n \n result = lst.copy()\n return result\n''' - in instrumented_source + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + decorated_original = original_code.replace( + "async def async_sorter", "@codeflash_behavior_async\nasync def async_sorter" ) + expected = sort_imports(code=inline_code + decorated_original, float_to_top=True) + assert instrumented_source.strip() == expected.strip() # Add codeflash capture instrument_codeflash_capture(func, {}, tests_root) @@ -300,10 +305,14 @@ async def test_async_perf(): # Verify the file was modified instrumented_source = fto_path.read_text("utf-8") - assert ( - instrumented_source - == '''import asyncio\nfrom typing import List, Union\n\nfrom codeflash.code_utils.codeflash_wrap_decorator import \\\n codeflash_performance_async\n\n\n@codeflash_performance_async\nasync def async_sorter(lst: List[Union[int, float]]) -> List[Union[int, float]]:\n """\n Async bubble sort implementation for testing.\n """\n print("codeflash stdout: Async sorting list")\n \n await asyncio.sleep(0.01)\n \n n = len(lst)\n for i in range(n):\n for j in range(0, n - i - 1):\n if lst[j] > lst[j + 1]:\n lst[j], lst[j + 1] = lst[j + 1], lst[j]\n \n result = lst.copy()\n print(f"result: {result}")\n return result\n\n\nclass AsyncBubbleSorter:\n """Class with async sorting method for testing."""\n \n async def sorter(self, lst: List[Union[int, float]]) -> List[Union[int, float]]:\n """\n Async bubble sort implementation within a class.\n """\n print("codeflash stdout: AsyncBubbleSorter.sorter() called")\n \n # Add some async delay\n await asyncio.sleep(0.005)\n \n n = len(lst)\n for i in range(n):\n for j in range(0, n - i - 1):\n if lst[j] > lst[j + 1]:\n lst[j], lst[j + 1] = lst[j + 1], lst[j]\n \n result = lst.copy()\n return result\n''' + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.PERFORMANCE) + decorated_original = original_code.replace( + "async def async_sorter", "@codeflash_performance_async\nasync def async_sorter" ) + expected = sort_imports(code=inline_code + decorated_original, float_to_top=True) + assert instrumented_source.strip() == expected.strip() instrument_codeflash_capture(func, {}, tests_root) @@ -411,61 +420,14 @@ async def async_error_function(lst): # Verify the file was modified instrumented_source = fto_path.read_text("utf-8") - expected_instrumented_source = """import asyncio -from typing import List, Union - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_behavior_async - - -async def async_sorter(lst: List[Union[int, float]]) -> List[Union[int, float]]: - \"\"\" - Async bubble sort implementation for testing. - \"\"\" - print("codeflash stdout: Async sorting list") - - await asyncio.sleep(0.01) - - n = len(lst) - for i in range(n): - for j in range(0, n - i - 1): - if lst[j] > lst[j + 1]: - lst[j], lst[j + 1] = lst[j + 1], lst[j] - - result = lst.copy() - print(f"result: {result}") - return result - + from codeflash.code_utils.formatter import sort_imports -class AsyncBubbleSorter: - \"\"\"Class with async sorting method for testing.\"\"\" - - async def sorter(self, lst: List[Union[int, float]]) -> List[Union[int, float]]: - \"\"\" - Async bubble sort implementation within a class. - \"\"\" - print("codeflash stdout: AsyncBubbleSorter.sorter() called") - - # Add some async delay - await asyncio.sleep(0.005) - - n = len(lst) - for i in range(n): - for j in range(0, n - i - 1): - if lst[j] > lst[j + 1]: - lst[j], lst[j + 1] = lst[j + 1], lst[j] - - result = lst.copy() - return result - - -@codeflash_behavior_async -async def async_error_function(lst): - \"\"\"Async function that raises an error for testing.\"\"\" - await asyncio.sleep(0.001) # Small delay - raise ValueError("Test error") -""" - assert expected_instrumented_source == instrumented_source + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + decorated_modified = modified_code.replace( + "async def async_error_function", "@codeflash_behavior_async\nasync def async_error_function" + ) + expected = sort_imports(code=inline_code + decorated_modified, float_to_top=True) + assert instrumented_source.strip() == expected.strip() instrument_codeflash_capture(func, {}, tests_root) opt = Optimizer( diff --git a/tests/test_instrument_async_tests.py b/tests/test_instrument_async_tests.py index 29e65ad06..d700d91d5 100644 --- a/tests/test_instrument_async_tests.py +++ b/tests/test_instrument_async_tests.py @@ -7,6 +7,7 @@ from codeflash.code_utils.instrument_existing_tests import ( add_async_decorator_to_function, + get_async_inline_code, inject_profiling_into_existing_test, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -57,20 +58,6 @@ def test_async_decorator_application_behavior_mode(temp_dir): async_function_code = ''' import asyncio -async def async_function(x: int, y: int) -> int: - """Simple async function for testing.""" - await asyncio.sleep(0.01) - return x * y -''' - - expected_decorated_code = ''' -import asyncio - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_behavior_async - - -@codeflash_behavior_async async def async_function(x: int, y: int) -> int: """Simple async function for testing.""" await asyncio.sleep(0.01) @@ -86,7 +73,15 @@ async def async_function(x: int, y: int) -> int: assert decorator_added modified_code = test_file.read_text() - assert modified_code.strip() == expected_decorated_code.strip() + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + expected = sort_imports( + code=inline_code + "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', + float_to_top=True, + ) + assert modified_code.strip() == expected.strip() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -94,20 +89,6 @@ def test_async_decorator_application_performance_mode(temp_dir): async_function_code = ''' import asyncio -async def async_function(x: int, y: int) -> int: - """Simple async function for testing.""" - await asyncio.sleep(0.01) - return x * y -''' - - expected_decorated_code = ''' -import asyncio - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_performance_async - - -@codeflash_performance_async async def async_function(x: int, y: int) -> int: """Simple async function for testing.""" await asyncio.sleep(0.01) @@ -123,7 +104,15 @@ async def async_function(x: int, y: int) -> int: assert decorator_added modified_code = test_file.read_text() - assert modified_code.strip() == expected_decorated_code.strip() + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.PERFORMANCE) + expected = sort_imports( + code=inline_code + "\n@codeflash_performance_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', + float_to_top=True, + ) + assert modified_code.strip() == expected.strip() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -132,20 +121,6 @@ def test_async_decorator_application_concurrency_mode(temp_dir): async_function_code = ''' import asyncio -async def async_function(x: int, y: int) -> int: - """Simple async function for testing.""" - await asyncio.sleep(0.01) - return x * y -''' - - expected_decorated_code = ''' -import asyncio - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_concurrency_async - - -@codeflash_concurrency_async async def async_function(x: int, y: int) -> int: """Simple async function for testing.""" await asyncio.sleep(0.01) @@ -161,7 +136,15 @@ async def async_function(x: int, y: int) -> int: assert decorator_added modified_code = test_file.read_text() - assert modified_code.strip() == expected_decorated_code.strip() + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.CONCURRENCY) + expected = sort_imports( + code=inline_code + "\n@codeflash_concurrency_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', + float_to_top=True, + ) + assert modified_code.strip() == expected.strip() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -182,27 +165,6 @@ def sync_method(self, a: int, b: int) -> int: return a - b ''' - expected_decorated_code = ''' -import asyncio - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_behavior_async - - -class Calculator: - """Test class with async methods.""" - - @codeflash_behavior_async - async def async_method(self, a: int, b: int) -> int: - """Async method in class.""" - await asyncio.sleep(0.005) - return a ** b - - def sync_method(self, a: int, b: int) -> int: - """Sync method in class.""" - return a - b -''' - test_file = temp_dir / "test_async.py" test_file.write_text(async_class_code) @@ -217,11 +179,31 @@ def sync_method(self, a: int, b: int) -> int: assert decorator_added modified_code = test_file.read_text() - assert modified_code.strip() == expected_decorated_code.strip() + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + expected = sort_imports( + code=inline_code + + "\nclass Calculator:\n" + ' """Test class with async methods."""\n' + " \n" + " @codeflash_behavior_async\n" + " async def async_method(self, a: int, b: int) -> int:\n" + ' """Async method in class."""\n' + " await asyncio.sleep(0.005)\n" + " return a ** b\n" + " \n" + " def sync_method(self, a: int, b: int) -> int:\n" + ' """Sync method in class."""\n' + " return a - b\n", + float_to_top=True, + ) + assert modified_code.strip() == expected.strip() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") def test_async_decorator_no_duplicate_application(temp_dir): + # Case 1: Old-style import already present — injector should detect and skip already_decorated_code = ''' from codeflash.code_utils.codeflash_wrap_decorator import codeflash_behavior_async import asyncio @@ -243,6 +225,30 @@ async def async_function(x: int, y: int) -> int: # Should not add duplicate decorator assert not decorator_added + # Case 2: Inline definition already present — injector should detect and skip + already_inline_code = ''' +import asyncio + +def codeflash_behavior_async(func): + return func + +@codeflash_behavior_async +async def async_function(x: int, y: int) -> int: + """Already decorated async function.""" + await asyncio.sleep(0.01) + return x * y +''' + + test_file2 = temp_dir / "test_async2.py" + test_file2.write_text(already_inline_code) + + func2 = FunctionToOptimize(function_name="async_function", file_path=test_file2, parents=[], is_async=True) + + decorator_added2 = add_async_decorator_to_function(test_file2, func2, TestingMode.BEHAVIOR) + + # Should not add duplicate decorator + assert not decorator_added2 + @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") def test_inject_profiling_async_function_behavior_mode(temp_dir): @@ -285,11 +291,17 @@ async def test_async_function(): assert source_success is True - # Verify the file was modified + # Verify the file was modified with exact expected output instrumented_source = source_file.read_text() - assert "@codeflash_behavior_async" in instrumented_source - assert "from codeflash.code_utils.codeflash_wrap_decorator import" in instrumented_source - assert "codeflash_behavior_async" in instrumented_source + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + expected = sort_imports( + code=inline_code + "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', + float_to_top=True, + ) + assert instrumented_source.strip() == expected.strip() success, instrumented_test_code = inject_profiling_into_existing_test( test_file, [CodePosition(8, 18), CodePosition(11, 19)], func, temp_dir, mode=TestingMode.BEHAVIOR @@ -340,12 +352,17 @@ async def test_async_function(): assert source_success is True - # Verify the file was modified + # Verify the file was modified with exact expected output instrumented_source = source_file.read_text() - assert "@codeflash_performance_async" in instrumented_source - # Check for the import with line continuation formatting - assert "from codeflash.code_utils.codeflash_wrap_decorator import" in instrumented_source - assert "codeflash_performance_async" in instrumented_source + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.PERFORMANCE) + expected = sort_imports( + code=inline_code + "\n@codeflash_performance_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', + float_to_top=True, + ) + assert instrumented_source.strip() == expected.strip() # Now test the full pipeline with source module path success, instrumented_test_code = inject_profiling_into_existing_test( @@ -406,11 +423,21 @@ async def test_mixed_functions(): # Verify the file was modified instrumented_source = source_file.read_text() - assert "@codeflash_behavior_async" in instrumented_source - assert "from codeflash.code_utils.codeflash_wrap_decorator import" in instrumented_source - assert "codeflash_behavior_async" in instrumented_source - # Sync function should remain unchanged - assert "def sync_function(x: int, y: int) -> int:" in instrumented_source + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + expected = sort_imports( + code=inline_code + + "\ndef sync_function(x: int, y: int) -> int:\n" + ' """Regular sync function."""\n' + " return x * y\n" + "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" + ' """Simple async function."""\n' + " await asyncio.sleep(0.01)\n" + " return x * y\n", + float_to_top=True, + ) + assert instrumented_source.strip() == expected.strip() success, instrumented_test_code = inject_profiling_into_existing_test( test_file, [CodePosition(8, 18), CodePosition(11, 19)], async_func, temp_dir, mode=TestingMode.BEHAVIOR @@ -446,24 +473,23 @@ async def nested_async_method(self, x: int) -> int: decorator_added = add_async_decorator_to_function(test_file, func, TestingMode.BEHAVIOR) - expected_output = """import asyncio - -from codeflash.code_utils.codeflash_wrap_decorator import \\ - codeflash_behavior_async - - -class OuterClass: - class InnerClass: - @codeflash_behavior_async - async def nested_async_method(self, x: int) -> int: - \"\"\"Nested async method.\"\"\" - await asyncio.sleep(0.001) - return x * 2 -""" - assert decorator_added modified_code = test_file.read_text() - assert modified_code.strip() == expected_output.strip() + from codeflash.code_utils.formatter import sort_imports + + inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + expected = sort_imports( + code=inline_code + + "\nclass OuterClass: \n" + " class InnerClass: \n" + " @codeflash_behavior_async\n" + " async def nested_async_method(self, x: int) -> int:\n" + ' """Nested async method."""\n' + " await asyncio.sleep(0.001)\n" + " return x * 2\n", + float_to_top=True, + ) + assert modified_code.strip() == expected.strip() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") From 64a18c9870f2c785f27ceb85abc32c1d9d8db129 Mon Sep 17 00:00:00 2001 From: KRRT7 Date: Wed, 18 Feb 2026 08:17:08 +0000 Subject: [PATCH 2/5] refactor: use helper file for async decorator instrumentation Replace inline code injection with a helper file approach that writes decorator implementations to a separate codeflash_async_wrapper.py file. This removes the codeflash package import dependency from instrumented source files while keeping line numbers stable (only 1 import + 1 decorator line added, same as before). Co-Authored-By: Claude Opus 4.6 --- .../code_directories/async_e2e/main.py | 6 +- .../code_utils/instrument_existing_tests.py | 81 +++++------- codeflash/optimization/function_optimizer.py | 25 +++- tests/test_async_run_and_parse_tests.py | 73 ++++++++--- tests/test_instrument_async_tests.py | 115 ++++++++---------- 5 files changed, 161 insertions(+), 139 deletions(-) diff --git a/code_to_optimize/code_directories/async_e2e/main.py b/code_to_optimize/code_directories/async_e2e/main.py index 317068a1c..8ab92ccdc 100644 --- a/code_to_optimize/code_directories/async_e2e/main.py +++ b/code_to_optimize/code_directories/async_e2e/main.py @@ -1,4 +1,3 @@ -import time import asyncio @@ -6,11 +5,14 @@ async def retry_with_backoff(func, max_retries=3): if max_retries < 1: raise ValueError("max_retries must be at least 1") last_exception = None + _sleep = asyncio.sleep for attempt in range(max_retries): try: return await func() except Exception as e: last_exception = e if attempt < max_retries - 1: - time.sleep(0.0001 * attempt) + delay = 0.0001 * attempt + if delay: + await _sleep(delay) raise last_exception diff --git a/codeflash/code_utils/instrument_existing_tests.py b/codeflash/code_utils/instrument_existing_tests.py index 8c53c8e01..cea455cf0 100644 --- a/codeflash/code_utils/instrument_existing_tests.py +++ b/codeflash/code_utils/instrument_existing_tests.py @@ -1700,69 +1700,44 @@ def get_async_inline_code(mode: TestingMode) -> str: return get_performance_async_inline_code() -class AsyncInlineCodeInjector(cst.CSTTransformer): - """Injects async decorator function definitions inline instead of importing from codeflash.""" +ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py" - def __init__(self, mode: TestingMode = TestingMode.BEHAVIOR) -> None: - self.mode = mode - self.has_inline_definition = False - self.has_old_import = False - - def _get_decorator_name(self) -> str: - if self.mode == TestingMode.BEHAVIOR: - return "codeflash_behavior_async" - if self.mode == TestingMode.CONCURRENCY: - return "codeflash_concurrency_async" - return "codeflash_performance_async" - - def visit_ImportFrom(self, node: cst.ImportFrom) -> None: - if ( - isinstance(node.module, cst.Attribute) - and isinstance(node.module.value, cst.Attribute) - and isinstance(node.module.value.value, cst.Name) - and node.module.value.value.value == "codeflash" - and node.module.value.attr.value == "code_utils" - and node.module.attr.value == "codeflash_wrap_decorator" - and not isinstance(node.names, cst.ImportStar) - ): - decorator_name = self._get_decorator_name() - for import_alias in node.names: - if import_alias.name.value == decorator_name: - self.has_old_import = True - def visit_FunctionDef(self, node: cst.FunctionDef) -> None: - if node.name.value == self._get_decorator_name(): - self.has_inline_definition = True +def get_decorator_name_for_mode(mode: TestingMode) -> str: + if mode == TestingMode.BEHAVIOR: + return "codeflash_behavior_async" + if mode == TestingMode.CONCURRENCY: + return "codeflash_concurrency_async" + return "codeflash_performance_async" + - def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: - if self.has_inline_definition or self.has_old_import: - return updated_node - inline_code = get_async_inline_code(self.mode) - inline_stmts = cst.parse_module(inline_code).body - return updated_node.with_changes(body=[*inline_stmts, *list(updated_node.body)]) +def write_async_helper_file(target_dir: Path, mode: TestingMode) -> Path: + """Write the async decorator helper file to the target directory.""" + helper_path = target_dir / ASYNC_HELPER_FILENAME + if helper_path.exists(): + decorator_name = get_decorator_name_for_mode(mode) + if f"def {decorator_name}" in helper_path.read_text("utf-8"): + return helper_path + helper_path.write_text(get_async_inline_code(mode), "utf-8") + return helper_path def add_async_decorator_to_function( - source_path: Path, function: FunctionToOptimize, mode: TestingMode = TestingMode.BEHAVIOR + source_path: Path, + function: FunctionToOptimize, + mode: TestingMode = TestingMode.BEHAVIOR, + project_root: Path | None = None, ) -> bool: """Add async decorator to an async function definition and write back to file. - Args: - ---- - source_path: Path to the source file to modify in-place. - function: The FunctionToOptimize object representing the target async function. - mode: The testing mode to determine which decorator to apply. - - Returns: - ------- - Boolean indicating whether the decorator was successfully added. + Writes a helper file containing the decorator implementation to project_root (or source directory + as fallback) and adds a standard import + decorator to the source file. """ if not function.is_async: return False try: - # Read source code with source_path.open(encoding="utf8") as f: source_code = f.read() @@ -1772,10 +1747,14 @@ def add_async_decorator_to_function( decorator_transformer = AsyncDecoratorAdder(function, mode) module = module.visit(decorator_transformer) - # Add the import if decorator was added if decorator_transformer.added_decorator: - import_transformer = AsyncInlineCodeInjector(mode) - module = module.visit(import_transformer) + # Write the helper file to project_root (on sys.path) or source dir as fallback + helper_dir = project_root if project_root is not None else source_path.parent + write_async_helper_file(helper_dir, mode) + # Add the import via CST so sort_imports can place it correctly + decorator_name = get_decorator_name_for_mode(mode) + import_node = cst.parse_statement(f"from codeflash_async_wrapper import {decorator_name}") + module = module.with_changes(body=[import_node, *list(module.body)]) modified_code = sort_imports(code=module.code, float_to_top=True) except Exception as e: diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index ed7f7f1fe..c4c68bf85 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -2296,7 +2296,10 @@ def establish_original_code_baseline( from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function success = add_async_decorator_to_function( - self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR + self.function_to_optimize.file_path, + self.function_to_optimize, + TestingMode.BEHAVIOR, + project_root=self.project_root, ) # Instrument codeflash capture @@ -2361,7 +2364,10 @@ def establish_original_code_baseline( from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function add_async_decorator_to_function( - self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE + self.function_to_optimize.file_path, + self.function_to_optimize, + TestingMode.PERFORMANCE, + project_root=self.project_root, ) try: @@ -2535,7 +2541,10 @@ def run_optimized_candidate( from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function add_async_decorator_to_function( - self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR + self.function_to_optimize.file_path, + self.function_to_optimize, + TestingMode.BEHAVIOR, + project_root=self.project_root, ) try: @@ -2611,7 +2620,10 @@ def run_optimized_candidate( from codeflash.code_utils.instrument_existing_tests import add_async_decorator_to_function add_async_decorator_to_function( - self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE + self.function_to_optimize.file_path, + self.function_to_optimize, + TestingMode.PERFORMANCE, + project_root=self.project_root, ) try: @@ -2974,7 +2986,10 @@ def run_concurrency_benchmark( try: # Add concurrency decorator to the source function add_async_decorator_to_function( - self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.CONCURRENCY + self.function_to_optimize.file_path, + self.function_to_optimize, + TestingMode.CONCURRENCY, + project_root=self.project_root, ) # Run the concurrency benchmark tests diff --git a/tests/test_async_run_and_parse_tests.py b/tests/test_async_run_and_parse_tests.py index 5750a015f..1777a1c73 100644 --- a/tests/test_async_run_and_parse_tests.py +++ b/tests/test_async_run_and_parse_tests.py @@ -8,8 +8,9 @@ import pytest from codeflash.code_utils.instrument_existing_tests import ( + ASYNC_HELPER_FILENAME, add_async_decorator_to_function, - get_async_inline_code, + get_decorator_name_for_mode, inject_profiling_into_existing_test, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -56,7 +57,9 @@ async def test_async_sort(): func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True) # For async functions, instrument the source module directly with decorators - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success @@ -64,11 +67,12 @@ async def test_async_sort(): instrumented_source = fto_path.read_text("utf-8") from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) decorated_original = original_code.replace( - "async def async_sorter", "@codeflash_behavior_async\nasync def async_sorter" + "async def async_sorter", f"@{decorator_name}\nasync def async_sorter" ) - expected = sort_imports(code=inline_code + decorated_original, float_to_top=True) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_original}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() # Add codeflash capture @@ -147,6 +151,9 @@ async def test_async_sort(): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -187,7 +194,9 @@ async def test_async_class_sort(): is_async=True, ) - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success @@ -269,6 +278,9 @@ async def test_async_class_sort(): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -299,7 +311,9 @@ async def test_async_perf(): func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True) # Instrument the source module with async performance decorators - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.PERFORMANCE) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.PERFORMANCE, project_root=project_root_path + ) assert source_success @@ -307,11 +321,12 @@ async def test_async_perf(): instrumented_source = fto_path.read_text("utf-8") from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.PERFORMANCE) + decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE) decorated_original = original_code.replace( - "async def async_sorter", "@codeflash_performance_async\nasync def async_sorter" + "async def async_sorter", f"@{decorator_name}\nasync def async_sorter" ) - expected = sort_imports(code=inline_code + decorated_original, float_to_top=True) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_original}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() instrument_codeflash_capture(func, {}, tests_root) @@ -368,6 +383,9 @@ async def test_async_perf(): # Clean up test files if test_path.exists(): test_path.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -413,7 +431,9 @@ async def async_error_function(lst): function_name="async_error_function", parents=[], file_path=Path(fto_path), is_async=True ) - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success @@ -422,11 +442,12 @@ async def async_error_function(lst): from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) decorated_modified = modified_code.replace( - "async def async_error_function", "@codeflash_behavior_async\nasync def async_error_function" + "async def async_error_function", f"@{decorator_name}\nasync def async_error_function" ) - expected = sort_imports(code=inline_code + decorated_modified, float_to_top=True) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{decorated_modified}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() instrument_codeflash_capture(func, {}, tests_root) @@ -488,6 +509,9 @@ async def async_error_function(lst): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -525,7 +549,9 @@ async def test_async_multi(): func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True) - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success instrument_codeflash_capture(func, {}, tests_root) @@ -598,6 +624,9 @@ async def test_async_multi(): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -640,7 +669,9 @@ async def test_async_edge_cases(): func = FunctionToOptimize(function_name="async_sorter", parents=[], file_path=Path(fto_path), is_async=True) - source_success = add_async_decorator_to_function(fto_path, func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + fto_path, func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success instrument_codeflash_capture(func, {}, tests_root) @@ -715,6 +746,9 @@ async def test_async_edge_cases(): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -949,7 +983,9 @@ async def test_mixed_sorting(): function_name="async_merge_sort", parents=[], file_path=Path(mixed_fto_path), is_async=True ) - source_success = add_async_decorator_to_function(mixed_fto_path, async_func, TestingMode.BEHAVIOR) + source_success = add_async_decorator_to_function( + mixed_fto_path, async_func, TestingMode.BEHAVIOR, project_root=project_root_path + ) assert source_success @@ -1022,3 +1058,6 @@ async def test_mixed_sorting(): test_path.unlink() if test_path_perf.exists(): test_path_perf.unlink() + helper_path = project_root_path / ASYNC_HELPER_FILENAME + if helper_path.exists(): + helper_path.unlink() diff --git a/tests/test_instrument_async_tests.py b/tests/test_instrument_async_tests.py index d700d91d5..0e57ec209 100644 --- a/tests/test_instrument_async_tests.py +++ b/tests/test_instrument_async_tests.py @@ -6,8 +6,9 @@ import pytest from codeflash.code_utils.instrument_existing_tests import ( + ASYNC_HELPER_FILENAME, add_async_decorator_to_function, - get_async_inline_code, + get_decorator_name_for_mode, inject_profiling_into_existing_test, ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -75,13 +76,14 @@ async def async_function(x: int, y: int) -> int: modified_code = test_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) - expected = sort_imports( - code=inline_code + "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) + code_with_decorator = async_function_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert modified_code.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -106,13 +108,14 @@ async def async_function(x: int, y: int) -> int: modified_code = test_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.PERFORMANCE) - expected = sort_imports( - code=inline_code + "\n@codeflash_performance_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE) + code_with_decorator = async_function_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert modified_code.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -138,13 +141,14 @@ async def async_function(x: int, y: int) -> int: modified_code = test_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.CONCURRENCY) - expected = sort_imports( - code=inline_code + "\n@codeflash_concurrency_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.CONCURRENCY) + code_with_decorator = async_function_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert modified_code.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -181,24 +185,14 @@ def sync_method(self, a: int, b: int) -> int: modified_code = test_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) - expected = sort_imports( - code=inline_code - + "\nclass Calculator:\n" - ' """Test class with async methods."""\n' - " \n" - " @codeflash_behavior_async\n" - " async def async_method(self, a: int, b: int) -> int:\n" - ' """Async method in class."""\n' - " await asyncio.sleep(0.005)\n" - " return a ** b\n" - " \n" - " def sync_method(self, a: int, b: int) -> int:\n" - ' """Sync method in class."""\n' - " return a - b\n", - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) + code_with_decorator = async_class_code.replace( + " async def async_method", f" @{decorator_name}\n async def async_method" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert modified_code.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") @@ -295,13 +289,14 @@ async def test_async_function(): instrumented_source = source_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) - expected = sort_imports( - code=inline_code + "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) + code_with_decorator = source_module_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() success, instrumented_test_code = inject_profiling_into_existing_test( test_file, [CodePosition(8, 18), CodePosition(11, 19)], func, temp_dir, mode=TestingMode.BEHAVIOR @@ -356,13 +351,14 @@ async def test_async_function(): instrumented_source = source_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.PERFORMANCE) - expected = sort_imports( - code=inline_code + "\n@codeflash_performance_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function for testing."""\n await asyncio.sleep(0.01)\n return x * y\n', - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.PERFORMANCE) + code_with_decorator = source_module_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() # Now test the full pipeline with source module path success, instrumented_test_code = inject_profiling_into_existing_test( @@ -425,19 +421,14 @@ async def test_mixed_functions(): instrumented_source = source_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) - expected = sort_imports( - code=inline_code - + "\ndef sync_function(x: int, y: int) -> int:\n" - ' """Regular sync function."""\n' - " return x * y\n" - "\n@codeflash_behavior_async\nasync def async_function(x: int, y: int) -> int:\n" - ' """Simple async function."""\n' - " await asyncio.sleep(0.01)\n" - " return x * y\n", - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) + code_with_decorator = source_module_code.replace( + "async def async_function", f"@{decorator_name}\nasync def async_function" ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert instrumented_source.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() success, instrumented_test_code = inject_profiling_into_existing_test( test_file, [CodePosition(8, 18), CodePosition(11, 19)], async_func, temp_dir, mode=TestingMode.BEHAVIOR @@ -477,19 +468,15 @@ async def nested_async_method(self, x: int) -> int: modified_code = test_file.read_text() from codeflash.code_utils.formatter import sort_imports - inline_code = get_async_inline_code(TestingMode.BEHAVIOR) - expected = sort_imports( - code=inline_code - + "\nclass OuterClass: \n" - " class InnerClass: \n" - " @codeflash_behavior_async\n" - " async def nested_async_method(self, x: int) -> int:\n" - ' """Nested async method."""\n' - " await asyncio.sleep(0.001)\n" - " return x * 2\n", - float_to_top=True, + decorator_name = get_decorator_name_for_mode(TestingMode.BEHAVIOR) + code_with_decorator = nested_async_code.replace( + " async def nested_async_method", + f" @{decorator_name}\n async def nested_async_method", ) + code_with_import = f"from codeflash_async_wrapper import {decorator_name}\n{code_with_decorator}" + expected = sort_imports(code=code_with_import, float_to_top=True) assert modified_code.strip() == expected.strip() + assert (temp_dir / ASYNC_HELPER_FILENAME).exists() @pytest.mark.skipif(sys.platform == "win32", reason="pending support for asyncio on windows") From e0861214ec08485361617ea7f83a0855f642db66 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Wed, 18 Feb 2026 03:46:02 -0500 Subject: [PATCH 3/5] fix: clean up async helper file and combine all decorators into single file Write all three async decorator implementations into one helper file to avoid overwrite issues when switching modes. Clean up the helper file in revert_code_and_helpers and early-exit paths so it doesn't persist in the user's project root after optimization. --- .../code_utils/instrument_existing_tests.py | 52 +++---------------- codeflash/optimization/function_optimizer.py | 9 ++++ 2 files changed, 15 insertions(+), 46 deletions(-) diff --git a/codeflash/code_utils/instrument_existing_tests.py b/codeflash/code_utils/instrument_existing_tests.py index cea455cf0..9486fc677 100644 --- a/codeflash/code_utils/instrument_existing_tests.py +++ b/codeflash/code_utils/instrument_existing_tests.py @@ -1497,11 +1497,11 @@ def _is_target_decorator(self, decorator_node: cst.Name | cst.Attribute | cst.Ca return False -def get_behavior_async_inline_code() -> str: - return """import asyncio +ASYNC_HELPER_INLINE_CODE = """import asyncio import gc import os import sqlite3 +import time from functools import wraps from pathlib import Path from tempfile import TemporaryDirectory @@ -1590,25 +1590,6 @@ async def async_wrapper(*args, **kwargs): raise exception return return_value return async_wrapper -""" - - -def get_performance_async_inline_code() -> str: - return """import asyncio -import gc -import os -from functools import wraps - - -def extract_test_context_from_env(): - test_module = os.environ["CODEFLASH_TEST_MODULE"] - test_class = os.environ.get("CODEFLASH_TEST_CLASS", None) - test_function = os.environ["CODEFLASH_TEST_FUNCTION"] - if test_module and test_function: - return (test_module, test_class if test_class else None, test_function) - raise RuntimeError( - "Test context environment variables not set - ensure tests are run through codeflash test runner" - ) def codeflash_performance_async(func): @@ -1649,15 +1630,6 @@ async def async_wrapper(*args, **kwargs): raise exception return return_value return async_wrapper -""" - - -def get_concurrency_async_inline_code() -> str: - return """import asyncio -import gc -import os -import time -from functools import wraps def codeflash_concurrency_async(func): @@ -1691,15 +1663,6 @@ async def async_wrapper(*args, **kwargs): return async_wrapper """ - -def get_async_inline_code(mode: TestingMode) -> str: - if mode == TestingMode.BEHAVIOR: - return get_behavior_async_inline_code() - if mode == TestingMode.CONCURRENCY: - return get_concurrency_async_inline_code() - return get_performance_async_inline_code() - - ASYNC_HELPER_FILENAME = "codeflash_async_wrapper.py" @@ -1711,14 +1674,11 @@ def get_decorator_name_for_mode(mode: TestingMode) -> str: return "codeflash_performance_async" -def write_async_helper_file(target_dir: Path, mode: TestingMode) -> Path: +def write_async_helper_file(target_dir: Path) -> Path: """Write the async decorator helper file to the target directory.""" helper_path = target_dir / ASYNC_HELPER_FILENAME - if helper_path.exists(): - decorator_name = get_decorator_name_for_mode(mode) - if f"def {decorator_name}" in helper_path.read_text("utf-8"): - return helper_path - helper_path.write_text(get_async_inline_code(mode), "utf-8") + if not helper_path.exists(): + helper_path.write_text(ASYNC_HELPER_INLINE_CODE, "utf-8") return helper_path @@ -1750,7 +1710,7 @@ def add_async_decorator_to_function( if decorator_transformer.added_decorator: # Write the helper file to project_root (on sys.path) or source dir as fallback helper_dir = project_root if project_root is not None else source_path.parent - write_async_helper_file(helper_dir, mode) + write_async_helper_file(helper_dir) # Add the import via CST so sort_imports can place it correctly decorator_name = get_decorator_name_for_mode(mode) import_node = cst.parse_statement(f"from codeflash_async_wrapper import {decorator_name}") diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index c4c68bf85..0a515076c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -1897,6 +1897,7 @@ def setup_and_establish_baseline( if self.args.override_fixtures: restore_conftest(original_conftest_content) cleanup_paths(paths_to_cleanup) + self.cleanup_async_helper_file() return Failure(baseline_result.failure()) original_code_baseline, test_functions_to_remove = baseline_result.unwrap() @@ -1908,6 +1909,7 @@ def setup_and_establish_baseline( if self.args.override_fixtures: restore_conftest(original_conftest_content) cleanup_paths(paths_to_cleanup) + self.cleanup_async_helper_file() return Failure("The threshold for test confidence was not met.") return Success( @@ -2279,6 +2281,13 @@ def revert_code_and_helpers(self, original_helper_code: dict[Path, str]) -> None self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) + self.cleanup_async_helper_file() + + def cleanup_async_helper_file(self) -> None: + from codeflash.code_utils.instrument_existing_tests import ASYNC_HELPER_FILENAME + + helper_path = self.project_root / ASYNC_HELPER_FILENAME + helper_path.unlink(missing_ok=True) def establish_original_code_baseline( self, From 950545119447c30dd79ca3414d3c0fe3ae603279 Mon Sep 17 00:00:00 2001 From: Kevin Turcios <106575910+KRRT7@users.noreply.github.com> Date: Wed, 18 Feb 2026 09:25:33 +0000 Subject: [PATCH 4/5] fix: revert async e2e fixture to use time.sleep() for optimization target The e2e test expects codeflash to detect and fix the intentional use of blocking time.sleep() in an async function. Using asyncio.sleep() removes the optimization opportunity and causes the CI job to fail. Co-Authored-By: Claude Opus 4.6 --- code_to_optimize/code_directories/async_e2e/main.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/code_to_optimize/code_directories/async_e2e/main.py b/code_to_optimize/code_directories/async_e2e/main.py index 8ab92ccdc..317068a1c 100644 --- a/code_to_optimize/code_directories/async_e2e/main.py +++ b/code_to_optimize/code_directories/async_e2e/main.py @@ -1,3 +1,4 @@ +import time import asyncio @@ -5,14 +6,11 @@ async def retry_with_backoff(func, max_retries=3): if max_retries < 1: raise ValueError("max_retries must be at least 1") last_exception = None - _sleep = asyncio.sleep for attempt in range(max_retries): try: return await func() except Exception as e: last_exception = e if attempt < max_retries - 1: - delay = 0.0001 * attempt - if delay: - await _sleep(delay) + time.sleep(0.0001 * attempt) raise last_exception From 6c092b5e7f73ccfc9ece386f2e7555c1ca8468dc Mon Sep 17 00:00:00 2001 From: Kevin Turcios <106575910+KRRT7@users.noreply.github.com> Date: Wed, 18 Feb 2026 09:45:51 +0000 Subject: [PATCH 5/5] fix: update expected coverage lines for optimized async e2e code The optimized code removes `import time`, shifting all function lines up by 1. Update expected_lines from [10-20] to [9-19] to match. Co-Authored-By: Claude Opus 4.6 --- tests/scripts/end_to_end_test_async.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/scripts/end_to_end_test_async.py b/tests/scripts/end_to_end_test_async.py index 0b4bf8957..0e38ae797 100644 --- a/tests/scripts/end_to_end_test_async.py +++ b/tests/scripts/end_to_end_test_async.py @@ -13,7 +13,7 @@ def run_test(expected_improvement_pct: int) -> bool: CoverageExpectation( function_name="retry_with_backoff", expected_coverage=100.0, - expected_lines=[10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20], + expected_lines=[9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19], ) ], )