enable async function optimization #769
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add is_async parameter to AiServiceClient.optimize_python_code method - Pass is_async flag to both optimization and test generation endpoints - Update optimization pipeline to pass is_async from FunctionToOptimize - Remove async function blocking check in function optimizer - Remove unused has_any_async_functions import This enables async functions to be properly processed by the optimization pipeline and sent to the AI service with the correct async context.
- Added comprehensive async test instrumentation (AsyncCallInstrumenter class) - Implemented async decorator functions (add_async_decorator_to_function, instrument_source_module_with_async_decorators) - Added async wrapper decorators (codeflash_behavior_async, codeflash_performance_async) - Updated edit_generated_tests.py to handle AsyncFunctionDef nodes in test parsing - Updated coverage_utils.py to include async functions in coverage analysis
- Add async throughput fields to Explanation dataclass - Implement throughput-based performance improvement calculation - Add MIN_THROUGHPUT_IMPROVEMENT_THRESHOLD configuration constant - Update explanation logic to prefer throughput metrics for async functions - Restore LSP compatibility with conditional test result display
in my case, i've started using symlinks a bit more often, and the current impl causes issues, we need to resolve the symlinked path too.
PR Reviewer Guide 🔍(Review updated until commit 6795a5c)Here are some key observations to aid the review process:
|
PR Code Suggestions ✨Latest suggestions up to 6795a5c
Previous suggestionsSuggestions up to commit 4325416
|
github-actions
bot
added
the
workflow-modified
⚡️ Codeflash found optimizations for this PR📄 24% (0.24x) speedup for
|
⚡️ Codeflash found optimizations for this PR📄 30% (0.30x) speedup for
|
The optimization achieves a 22% speedup by eliminating redundant regex compilation and reducing unnecessary string operations. **Key optimizations:** 1. **Pre-compiled regex patterns**: The original code compiled the same regex pattern multiple times (3,114 compilations taking 43.4% of total time). The optimized version compiles each pattern only once upfront using `_compile_function_patterns()`, moving this expensive operation outside the nested loops. 2. **Efficient string manipulation**: Instead of using `re.sub()` which searches the entire string again, the optimized version uses `finditer()` to get match positions directly, then performs string slicing (`source[:start] + source[end:]`) to remove matched functions. This avoids the overhead of regex substitution. 3. **Early termination**: After finding and removing a function match, the code breaks from the inner loop since only one match per function is expected, preventing unnecessary continued iteration. **Performance impact by test case:** - The optimizations are most effective for scenarios with multiple test functions to remove across multiple generated tests (the typical use case) - For edge cases like empty test lists, there's minimal overhead from pre-compilation but no significant benefit - The approach maintains correct behavior for decorated functions (skipping `@pytest.mark.parametrize` functions as intended) The line profiler shows the regex compilation time dropped from 43.4% to being absorbed into the 89.8% upfront compilation cost, while the substitution overhead (51.7% in original) is eliminated entirely.
⚡️ Codeflash found optimizations for this PR📄 23% (0.23x) speedup for
|
…25-09-27T01.33.17 ⚡️ Speed up function `remove_functions_from_generated_tests` by 23% in PR #769 (`clean-async-branch`)
|
This PR is now faster! 🚀 @KRRT7 accepted my optimizations from: |
The optimized code achieves a **123% speedup** by replacing expensive AST traversal operations with more efficient alternatives: **Key Optimizations:** 1. **Decorator Search Optimization**: Replaced the `any()` generator expression with a simple loop that breaks early when finding `timeout_decorator.timeout`. This avoids unnecessary attribute lookups and iterations through the decorator list, especially beneficial when the decorator is found early or when there are many decorators. 2. **AST Traversal Replacement**: The most significant optimization replaces `ast.walk(stmt)` with a manual stack-based depth-first search in `_optimized_instrument_statement()`. The original `ast.walk()` creates a list of every node in the AST subtree, which is memory-intensive and includes many irrelevant nodes. The optimized version: - Uses a stack to traverse nodes manually - Only explores child nodes via `_fields` attribute access - Immediately returns when finding an `ast.Await` node that matches criteria - Avoids creating intermediate collections **Performance Impact by Test Case:** - **Large-scale tests** see the biggest improvements (125-129% faster) because they have many statements to traverse - **Nested structures** benefit significantly (57-93% faster) as the optimization avoids deep, unnecessary traversals - **Simple test cases** still see 29-48% improvements from the decorator optimization - **Functions with many await calls** show excellent scaling (123-127% faster) due to reduced per-statement traversal costs The line profiler shows the critical bottleneck was in `_instrument_statement()` (96.4% of time originally), which is now reduced to 93.3% but with much lower absolute time, demonstrating the effectiveness of the AST traversal optimization.
⚡️ Codeflash found optimizations for this PR📄 123% (1.23x) speedup for
|
⚡️ Codeflash found optimizations for this PR📄 39% (0.39x) speedup for
|
…25-09-27T02.50.03 ⚡️ Speed up method `AsyncCallInstrumenter.visit_AsyncFunctionDef` by 123% in PR #769 (`clean-async-branch`)
|
This PR is now faster! 🚀 @KRRT7 accepted my optimizations from: |
misrasaurabh1
previously approved these changes
Sep 29, 2025
![codeflash-ai[bot]](https://avatars.githubusercontent.com/in/413177?s=60&v=4){kind=small}
Comment on lines
+325
to
+335
| if ( | ||
| isinstance(item, ast.FunctionDef) | ||
| and item.name.startswith("test_") | ||
| and not any( | ||
| isinstance(d, ast.Call) | ||
| and isinstance(d.func, ast.Name) | ||
| and d.func.id == "timeout_decorator.timeout" | ||
| for d in item.decorator_list | ||
| ) | ||
| ): | ||
| item.decorator_list.append(timeout_decorator) |
There was a problem hiding this comment.
⚡️Codeflash found 10% (0.10x) speedup for AsyncCallInstrumenter.visit_ClassDef in codeflash/code_utils/instrument_existing_tests.py
⏱️ Runtime : 4.11 milliseconds → 3.72 milliseconds (best of 216 runs)
📝 Explanation and details
The optimization replaces the computationally expensive any() generator expression with an explicit loop and early break pattern. In the original code, any() creates a generator that evaluates a complex condition for each decorator in item.decorator_list, even when an early match is found. The optimized version uses a simple flag-based approach with has_timeout_decorator = False and an explicit loop that breaks immediately upon finding the target decorator.
This change eliminates the overhead of:
- Generator function creation and iteration protocol
- Repeated evaluation of the complex isinstance() chain for all decorators when early matches occur
- The
any()builtin call overhead
The optimization is particularly effective for test cases with many test methods (13.5% speedup on 500 test functions) because it reduces the per-method decorator checking cost. For smaller classes, the speedup is more modest (5-11%) but still consistent. The explicit loop pattern allows the CPU to exit the decorator search as soon as a timeout decorator is found, rather than continuing to evaluate the generator expression.
✅ Correctness verification report:
| Test | Status |
|---|---|
| ⚙️ Existing Unit Tests | 🔘 None Found |
| 🌀 Generated Regression Tests | ✅ 50 Passed |
| ⏪ Replay Tests | 🔘 None Found |
| 🔎 Concolic Coverage Tests | 🔘 None Found |
| 📊 Tests Coverage | 100.0% |
🌀 Generated Regression Tests and Runtime
from __future__ import annotations
import ast
# imports
import pytest # used for our unit tests
from codeflash.code_utils.instrument_existing_tests import \
AsyncCallInstrumenter
# Dummy classes for dependencies
class FunctionToOptimize:
def __init__(self, function_name, parents=None, top_level_parent_name=None):
self.function_name = function_name
self.parents = parents or []
self.top_level_parent_name = top_level_parent_name
class CodePosition:
pass
class TestingMode:
BEHAVIOR = "behavior"
OTHER = "other"
from codeflash.code_utils.instrument_existing_tests import \
AsyncCallInstrumenter
# Helper to parse code and return AST
def parse_classdef(code: str) -> ast.ClassDef:
mod = ast.parse(code)
for node in mod.body:
if isinstance(node, ast.ClassDef):
return node
raise ValueError("No ClassDef found in code")
# Helper to check if a function has the timeout decorator
def has_timeout_decorator(func: ast.FunctionDef) -> bool:
for dec in func.decorator_list:
if (
isinstance(dec, ast.Call)
and isinstance(dec.func, ast.Name)
and dec.func.id == "timeout_decorator.timeout"
):
return True
return False
# ---------------- BASIC TEST CASES ----------------
def test_no_test_methods_unittest():
# Class with no test methods: nothing should be decorated
code = """
class MyTestCase:
def helper(self): pass
def not_a_test(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 10.2μs -> 10.2μs (0.098% slower)
# No function should have the timeout decorator
for item in result.body:
if isinstance(item, ast.FunctionDef):
pass
def test_single_test_method_unittest():
# Class with one test method: should be decorated
code = """
class MyTestCase:
def test_simple(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 12.2μs -> 11.6μs (5.90% faster)
# The test_simple method should have the timeout decorator
for item in result.body:
if isinstance(item, ast.FunctionDef) and item.name == "test_simple":
pass
def test_multiple_test_methods_unittest():
# Class with multiple test methods: all should be decorated
code = """
class MyTestCase:
def test_one(self): pass
def test_two(self): pass
def helper(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 15.2μs -> 14.2μs (7.18% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef):
if item.name.startswith("test_"):
pass
else:
pass
def test_non_unittest_framework():
# Should not decorate anything if not using unittest
code = """
class MyTestCase:
def test_one(self): pass
def test_two(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"pytest", # not unittest
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 9.34μs -> 9.33μs (0.096% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef):
pass
def test_already_has_timeout_decorator():
# Should not add duplicate timeout decorator
code = """
class MyTestCase:
@timeout_decorator.timeout(15)
def test_one(self): pass
def test_two(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 14.9μs -> 13.6μs (9.52% faster)
found = 0
for item in result.body:
if isinstance(item, ast.FunctionDef):
if item.name == "test_one":
found += 1
elif item.name == "test_two":
found += 1
# ---------------- EDGE TEST CASES ----------------
def test_decorator_list_is_empty():
# Handles function with no decorators
code = """
class MyTestCase:
def test_empty_decorator(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 11.3μs -> 10.6μs (7.12% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef) and item.name == "test_empty_decorator":
pass
def test_non_function_items_in_class():
# Handles classes with non-function items
code = """
class MyTestCase:
x = 5
def test_x(self): pass
@property
def y(self): return 10
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 23.0μs -> 22.0μs (4.69% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef):
if item.name == "test_x":
pass
elif item.name == "y":
pass
def test_function_name_edge_cases():
# Only functions starting with test_ get decorated
code = """
class MyTestCase:
def test(self): pass
def testCase(self): pass
def test_abc(self): pass
def test123(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 14.3μs -> 13.4μs (6.51% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef):
if item.name == "test_abc":
pass
else:
pass
def test_decorator_is_not_call():
# Handles decorators that are not ast.Call
code = """
class MyTestCase:
@staticmethod
def test_static(self): pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 11.9μs -> 10.8μs (9.88% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef) and item.name == "test_static":
# Should have both staticmethod and timeout_decorator
found_static = any(isinstance(dec, ast.Name) and dec.id == "staticmethod" for dec in item.decorator_list)
def test_empty_class_body():
# Class with no body should not raise
code = """
class MyTestCase:
pass
"""
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
# Should not raise
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 8.12μs -> 8.03μs (1.12% faster)
# ---------------- LARGE SCALE TEST CASES ----------------
def test_large_number_of_test_methods():
# Large class with many test methods
N = 500
code = "class BigTestCase:\n" + "\n".join(
f" def test_{i}(self): pass" for i in range(N)
)
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 954μs -> 840μs (13.6% faster)
count = 0
for item in result.body:
if isinstance(item, ast.FunctionDef):
count += 1
def test_large_class_mixed_methods():
# Large class with mixed test and non-test methods
N = 250
code = "class MixedTestCase:\n" + "\n".join(
f" def test_{i}(self): pass" for i in range(N)
) + "\n" + "\n".join(
f" def helper_{i}(self): pass" for i in range(N)
)
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"unittest",
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 641μs -> 569μs (12.5% faster)
test_count = helper_count = 0
for item in result.body:
if isinstance(item, ast.FunctionDef):
if item.name.startswith("test_"):
test_count += 1
else:
helper_count += 1
def test_large_non_unittest():
# Large class with unittest-like methods but non-unittest framework
N = 300
code = "class BigTestCase:\n" + "\n".join(
f" def test_{i}(self): pass" for i in range(N)
)
classdef = parse_classdef(code)
instrumenter = AsyncCallInstrumenter(
FunctionToOptimize("foo"),
"mod.py",
"pytest", # not unittest
[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); result = codeflash_output # 367μs -> 356μs (3.08% faster)
for item in result.body:
if isinstance(item, ast.FunctionDef):
pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations
import ast
# imports
import pytest # used for our unit tests
from codeflash.code_utils.instrument_existing_tests import \
AsyncCallInstrumenter
class DummyFunction:
def __init__(self, function_name, parents=None, top_level_parent_name=None):
self.function_name = function_name
self.parents = parents or []
self.top_level_parent_name = top_level_parent_name
# Helper to parse code and get ClassDef node
def get_classdef_node(source: str) -> ast.ClassDef:
tree = ast.parse(source)
for node in tree.body:
if isinstance(node, ast.ClassDef):
return node
raise ValueError("No ClassDef found in source")
# Helper to check if a functiondef has the timeout decorator
def has_timeout_decorator(funcdef: ast.FunctionDef) -> bool:
for d in funcdef.decorator_list:
if (
isinstance(d, ast.Call)
and isinstance(d.func, ast.Name)
and d.func.id == "timeout_decorator.timeout"
and len(d.args) == 1
and isinstance(d.args[0], ast.Constant)
and d.args[0].value == 15
):
return True
return False
# Basic Test Cases
def test_unittest_adds_timeout_to_test_functions():
# Test that the decorator is added to test functions for unittest
src = '''
class TestFoo:
def test_one(self): pass
def test_two(self): pass
def helper(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_one"),
module_path="foo.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 16.1μs -> 15.1μs (6.38% faster)
test_funcs = [item for item in new_node.body if isinstance(item, ast.FunctionDef)]
# Only test_ functions should get the decorator
for func in test_funcs:
if func.name.startswith("test_"):
pass
else:
pass
def test_pytest_does_not_add_timeout():
# Test that no decorator is added for pytest
src = '''
class TestBar:
def test_a(self): pass
def test_b(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_a"),
module_path="bar.py",
test_framework="pytest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 9.48μs -> 9.42μs (0.637% faster)
for func in [item for item in new_node.body if isinstance(item, ast.FunctionDef)]:
pass
def test_no_test_functions():
# Class with no test_ functions should not get any decorator
src = '''
class HelperClass:
def foo(self): pass
def bar(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("foo"),
module_path="helper.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 9.62μs -> 9.32μs (3.22% faster)
for func in [item for item in new_node.body if isinstance(item, ast.FunctionDef)]:
pass
def test_existing_timeout_decorator_not_duplicated():
# If timeout_decorator.timeout is already present, do not duplicate
src = '''
class TestDup:
@timeout_decorator.timeout(15)
def test_dup(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_dup"),
module_path="dup.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 12.1μs -> 11.1μs (8.65% faster)
func = [item for item in new_node.body if isinstance(item, ast.FunctionDef)][0]
# Edge Test Cases
def test_class_with_non_function_body_items():
# Class with assignments and other nodes in body
src = '''
class TestEdge:
x = 1
def test_func(self): pass
y = 2
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_func"),
module_path="edge.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 28.1μs -> 26.7μs (5.18% faster)
# Only test_func should get the decorator
for item in new_node.body:
if isinstance(item, ast.FunctionDef):
if item.name == "test_func":
pass
else:
pass
else:
pass
def test_class_with_decorators_on_test_function():
# Test function with other decorators should still get timeout_decorator
src = '''
class TestDecorators:
@other_decorator
def test_deco(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_deco"),
module_path="deco.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 12.1μs -> 10.9μs (11.1% faster)
func = [item for item in new_node.body if isinstance(item, ast.FunctionDef)][0]
def test_class_with_inherited_test_functions():
# Class with test functions that are inherited (should not be touched)
src = '''
class Base:
def test_base(self): pass
class Derived(Base):
def test_derived(self): pass
'''
tree = ast.parse(src)
# Only Derived should be instrumented
derived = [n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "Derived"][0]
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_derived"),
module_path="inherit.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(derived); new_node = codeflash_output # 14.9μs -> 13.7μs (8.64% faster)
func = [item for item in new_node.body if isinstance(item, ast.FunctionDef)][0]
def test_class_with_test_function_named_exactly_test():
# Function named 'test' (not 'test_') should not get decorator
src = '''
class TestExact:
def test(self): pass
def test_one(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_one"),
module_path="exact.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 12.6μs -> 11.8μs (6.43% faster)
for func in [item for item in new_node.body if isinstance(item, ast.FunctionDef)]:
if func.name == "test_one":
pass
else:
pass
def test_class_with_async_test_functions():
# Async test functions should also be instrumented
src = '''
class TestAsync:
async def test_async(self): pass
'''
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_async"),
module_path="async.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 12.2μs -> 12.1μs (0.577% faster)
func = [item for item in new_node.body if isinstance(item, ast.AsyncFunctionDef)][0]
# Large Scale Test Cases
def test_large_number_of_test_functions():
# Class with many test_ functions
NUM_FUNCS = 500
src_lines = ["class TestLarge:"]
for i in range(NUM_FUNCS):
src_lines.append(f" def test_func_{i}(self): pass")
src = "\n".join(src_lines)
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_func_0"),
module_path="large.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 959μs -> 845μs (13.5% faster)
test_funcs = [item for item in new_node.body if isinstance(item, ast.FunctionDef)]
for func in test_funcs:
pass
def test_large_number_of_non_test_functions():
# Class with many non-test functions
NUM_FUNCS = 500
src_lines = ["class HelperLarge:"]
for i in range(NUM_FUNCS):
src_lines.append(f" def helper_func_{i}(self): pass")
src = "\n".join(src_lines)
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("helper_func_0"),
module_path="largehelper.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 299μs -> 294μs (1.85% faster)
helper_funcs = [item for item in new_node.body if isinstance(item, ast.FunctionDef)]
for func in helper_funcs:
pass
def test_large_mixed_class():
# Class with both test_ and non-test functions
NUM_TEST_FUNCS = 250
NUM_HELPER_FUNCS = 250
src_lines = ["class MixedLarge:"]
for i in range(NUM_TEST_FUNCS):
src_lines.append(f" def test_func_{i}(self): pass")
for i in range(NUM_HELPER_FUNCS):
src_lines.append(f" def helper_func_{i}(self): pass")
src = "\n".join(src_lines)
classdef = get_classdef_node(src)
instrumenter = AsyncCallInstrumenter(
DummyFunction("test_func_0"),
module_path="mixedlarge.py",
test_framework="unittest",
call_positions=[],
)
codeflash_output = instrumenter.visit_ClassDef(classdef); new_node = codeflash_output # 631μs -> 573μs (10.1% faster)
funcs = [item for item in new_node.body if isinstance(item, ast.FunctionDef)]
for func in funcs:
if func.name.startswith("test_"):
pass
else:
pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.To test or edit this optimization locally git merge codeflash/optimize-pr769-2025-09-29T21.00.36
Click to see suggested changes
Suggested change
| if ( | |
| isinstance(item, ast.FunctionDef) | |
| and item.name.startswith("test_") | |
| and not any( | |
| isinstance(d, ast.Call) | |
| and isinstance(d.func, ast.Name) | |
| and d.func.id == "timeout_decorator.timeout" | |
| for d in item.decorator_list | |
| ) | |
| ): | |
| item.decorator_list.append(timeout_decorator) | |
| if isinstance(item, ast.FunctionDef) and item.name.startswith("test_"): | |
| has_timeout_decorator = False | |
| for d in item.decorator_list: | |
| if ( | |
| isinstance(d, ast.Call) | |
| and isinstance(d.func, ast.Name) | |
| and d.func.id == "timeout_decorator.timeout" | |
| ): | |
| has_timeout_decorator = True | |
| break | |
| if not has_timeout_decorator: | |
| item.decorator_list.append(timeout_decorator) |
There was a problem hiding this comment.
LGTM - but I'll merge this in later.
aseembits93
approved these changes
Sep 29, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
User description
less messy
PR Type
Enhancement, Tests
Description
Add comprehensive async test suite
Enable async optimization pipeline end-to-end
Throughput-aware performance selection logic
Fix coverage utils empty-data handling
Diagram Walkthrough
File Walkthrough
3 files
End-to-end async run/parse tests and behaviorsAsync decorator injection and profiling testsAsync helpers detection and revert scenarios2 files
Async-aware pipeline, benchmarking, throughput, explanationsAdd throughput gain and async-aware speedup logic1 files
Fix empty coverage return type and checks29 files