Skip to content

support async functions in codeflash #617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 25 commits into from
Closed

support async functions in codeflash #617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
1abf574
console rulez
KRRT7 Aug 1, 2025
8654ed8
Update functions_to_optimize.py
KRRT7 Aug 1, 2025
bd59fe7
extract async topkek
KRRT7 Aug 1, 2025
4112eff
Update function_optimizer.py
KRRT7 Aug 1, 2025
557d1ab
Update coverage_utils.py
KRRT7 Aug 1, 2025
479da54
Update sentry.py
KRRT7 Aug 1, 2025
c07d9d1
discover async functions properly
KRRT7 Aug 5, 2025
d948ad7
first pass async wrapper
KRRT7 Aug 5, 2025
4132759
revert changes
KRRT7 Aug 6, 2025
f1de751
add async codeflash wrap string
KRRT7 Aug 6, 2025
7ec754c
another impl
KRRT7 Aug 6, 2025
42fb4f0
temp
KRRT7 Aug 6, 2025
b957d3c
Update instrument_existing_tests.py
KRRT7 Aug 6, 2025
b02a5bb
Update instrument_existing_tests.py
KRRT7 Aug 6, 2025
eabfd2b
Update uv.lock
KRRT7 Aug 6, 2025
0042d08
async tests instrumentation
KRRT7 Aug 16, 2025
f61c139
revert changes
KRRT7 Aug 16, 2025
99b4300
pre-commit
KRRT7 Aug 16, 2025
3bafc8d
Update instrument_existing_tests.py
KRRT7 Aug 16, 2025
a150e43
uv lock refresh
KRRT7 Aug 16, 2025
f2c0740
libcst at it again
KRRT7 Aug 16, 2025
66eac2a
just bool it
KRRT7 Aug 16, 2025
e292c5e
re-revert
KRRT7 Aug 16, 2025
a5182c6
Merge branch 'main' into alpha-async
KRRT7 Aug 18, 2025
c4e3e00
Merge branch 'main' into alpha-async
KRRT7 Sep 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions code_to_optimize/async_adder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
import asyncio


async def async_add(a, b):
"""Simple async function that adds two numbers."""
await asyncio.sleep(0.001) # Simulate some async work
print(f"codeflash stdout: Adding {a} + {b}")
result = a + b
print(f"result: {result}")
return result
7 changes: 7 additions & 0 deletions code_to_optimize/code_directories/async_repo/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
[tool.codeflash]
# All paths are relative to this pyproject.toml's directory.
module-root = "."
tests-root = "tests"
test-framework = "pytest"
ignore-paths = []
formatter-cmds = ["black $file"]
5 changes: 5 additions & 0 deletions code_to_optimize/code_directories/async_repo/shocker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
from time import sleep

async def tasked():
sleep(0.002)
return "Tasked"
26 changes: 26 additions & 0 deletions code_to_optimize/code_directories/async_repo/tests/test_shocker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
import pytest
import asyncio
from shocker import tasked


@pytest.mark.asyncio
async def test_tasked_basic():
result = await tasked()
assert result == "Tasked"


@pytest.mark.asyncio
async def test_tasked_gather():
results = await asyncio.gather(*(tasked() for _ in range(5)))
assert results == ["Tasked"] * 5


def test_tasked_many_parallel_invocations():
async def run_many():
tasks = [tasked() for _ in range(1000)]
results = await asyncio.gather(*tasks)
return results

results = asyncio.run(run_many())
assert len(results) == 1000, "Should return 1000 results"
assert all(r == "Tasked" for r in results), "All results should be 'Tasked'"
226 changes: 214 additions & 12 deletions codeflash/code_utils/instrument_existing_tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,39 @@ def is_argument_name(name: str, arguments_node: ast.arguments) -> bool:
)


class AsyncIOGatherRemover(ast.NodeTransformer):
def _contains_asyncio_gather(self, node: ast.AST) -> bool:
"""Check if a node contains asyncio.gather calls."""
for child_node in ast.walk(node):
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Attribute)
and isinstance(child_node.func.value, ast.Name)
and child_node.func.value.id == "asyncio"
and child_node.func.attr == "gather"
):
return True

if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Name)
and child_node.func.id == "gather"
):
return True

Comment on lines +54 to +70
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 71% (0.71x) speedup for AsyncIOGatherRemover._contains_asyncio_gather in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 23.2 milliseconds → 13.6 milliseconds (best of 73 runs)

📝 Explanation and details

The optimization replaces ast.walk() with a manual stack-based traversal that enables early termination. Here's why this delivers a 70% speedup:

Key optimization: Early return prevents unnecessary traversal

  • Original: ast.walk() always visits every node in the AST, even after finding a match
  • Optimized: Manual stack stops traversal immediately when asyncio.gather or gather is found

Performance benefits by test case type:

  • Small ASTs with early matches (simple gather calls): 74-80% faster - the optimization finds the target quickly and exits
  • Large ASTs with one gather call: Up to 6104% faster - avoids walking through hundreds of unnecessary nodes after finding the match
  • Large ASTs with no gather: ~6-8% faster - still must traverse all nodes, but stack operations are slightly more efficient than ast.walk()
  • Large ASTs with gather at start: Dramatic speedup as traversal stops immediately

Technical improvements:

  • Uses ast.iter_child_nodes() with explicit stack instead of generator-based ast.walk()
  • Consolidates the two separate condition checks into a single flow
  • Reduces redundant isinstance calls by storing child_node.func in a variable

The optimization is most effective when gather calls exist in the code, as it can terminate search early rather than examining the entire AST structure.

✅ Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 134 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 AsyncIOGatherRemover

# unit tests

@pytest.fixture
def remover():
    """Fixture to provide an instance of AsyncIOGatherRemover."""
    return AsyncIOGatherRemover()

# ----------------------
# Basic Test Cases
# ----------------------

def test_no_calls(remover):
    """Test a function with no function calls at all."""
    code = "def foo():\n    x = 1\n    return x"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 16.1μs -> 14.9μs (7.87% faster)

def test_simple_asyncio_gather_call(remover):
    """Test a direct call to asyncio.gather."""
    code = "async def foo():\n    await asyncio.gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 14.5μs -> 8.34μs (74.2% faster)

def test_simple_gather_call(remover):
    """Test a direct call to gather (unqualified)."""
    code = "async def foo():\n    await gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 14.0μs -> 7.92μs (76.6% faster)

def test_other_asyncio_method(remover):
    """Test a call to another asyncio method (not gather)."""
    code = "async def foo():\n    await asyncio.sleep(1)"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 16.7μs -> 15.4μs (8.59% faster)

def test_gather_as_variable(remover):
    """Test where 'gather' is used as a variable, not a function."""
    code = "def foo():\n    gather = 5\n    return gather"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 15.2μs -> 14.2μs (6.89% faster)

def test_gather_as_function_name(remover):
    """Test where 'gather' is the name of a function, not a call."""
    code = "def gather():\n    return 1"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 11.3μs -> 9.43μs (19.8% faster)

def test_asyncio_gather_as_variable(remover):
    """Test where 'asyncio.gather' is used as a variable, not a call."""
    code = "def foo():\n    x = asyncio.gather\n    return x"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 16.8μs -> 15.7μs (6.94% faster)

# ----------------------
# Edge Test Cases
# ----------------------

def test_nested_gather_call(remover):
    """Test gather call nested inside another function call."""
    code = "def foo():\n    bar(gather(a(), b()))"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.5μs -> 8.66μs (78.8% faster)

def test_nested_asyncio_gather_call(remover):
    """Test asyncio.gather call nested inside another function call."""
    code = "def foo():\n    bar(asyncio.gather(a(), b()))"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.5μs -> 8.67μs (79.1% faster)

def test_gather_in_lambda(remover):
    """Test gather call inside a lambda."""
    code = "f = lambda: gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.6μs -> 6.40μs (113% faster)

def test_asyncio_gather_in_comprehension(remover):
    """Test asyncio.gather call inside a comprehension."""
    code = "results = [asyncio.gather(x()) for x in funcs]"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.0μs -> 11.5μs (4.55% faster)

def test_gather_in_class_method(remover):
    """Test gather call inside a class method."""
    code = "class A:\n    def foo(self):\n        return gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.6μs -> 8.67μs (80.0% faster)

def test_gather_in_decorator(remover):
    """Test gather call used as a decorator."""
    code = "@gather\ndef foo():\n    pass"
    tree = ast.parse(code)
    # This is not a function call, just decorator syntax, so should be False
    codeflash_output = not remover._contains_asyncio_gather(tree) # 11.6μs -> 10.1μs (14.7% faster)

def test_gather_shadowed_by_import(remover):
    """Test gather shadowed by import from another module."""
    code = "from mymod import gather\ndef foo():\n    gather(a, b)"
    tree = ast.parse(code)
    # The AST can't know if this is asyncio.gather, but it's still a call to 'gather'
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.3μs -> 6.68μs (129% faster)

def test_asyncio_gather_shadowed_by_assignment(remover):
    """Test 'asyncio.gather' is shadowed by assignment."""
    code = "asyncio.gather = lambda *a: None\ndef foo():\n    asyncio.gather(a(), b())"
    tree = ast.parse(code)
    # The AST can't know about shadowing, so it should still return True
    codeflash_output = remover._contains_asyncio_gather(tree) # 20.2μs -> 6.86μs (195% faster)

def test_gather_with_different_case(remover):
    """Test 'Gather' with different case -- should not match."""
    code = "def foo():\n    Gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 17.9μs -> 17.4μs (2.52% faster)

def test_asyncio_gather_with_different_case(remover):
    """Test 'asyncio.Gather' with different case -- should not match."""
    code = "def foo():\n    asyncio.Gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 19.6μs -> 18.7μs (5.21% faster)

def test_gather_in_try_except(remover):
    """Test gather call inside try/except."""
    code = "try:\n    gather(a(), b())\nexcept Exception:\n    pass"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.5μs -> 10.1μs (23.6% faster)

def test_gather_in_if_else(remover):
    """Test gather call inside an if/else block."""
    code = "if cond:\n    gather(a(), b())\nelse:\n    pass"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.6μs -> 6.96μs (80.7% faster)

def test_gather_in_multiple_functions(remover):
    """Test gather call present in only one of several functions."""
    code = "def foo():\n    pass\ndef bar():\n    gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 16.6μs -> 6.66μs (149% faster)

def test_asyncio_gather_in_async_function(remover):
    """Test asyncio.gather call inside an async function."""
    code = "async def foo():\n    await asyncio.gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.8μs -> 7.93μs (74.2% faster)

def test_gather_in_with_statement(remover):
    """Test gather call inside a with statement."""
    code = "with open('f') as f:\n    gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.3μs -> 6.35μs (110% faster)

def test_gather_in_while_loop(remover):
    """Test gather call inside a while loop."""
    code = "while True:\n    gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 11.2μs -> 6.18μs (81.8% faster)

def test_gather_in_augmented_assignment(remover):
    """Test gather call as part of an augmented assignment."""
    code = "x += gather(a(), b())"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 10.8μs -> 4.82μs (124% faster)

def test_gather_in_list(remover):
    """Test gather call as an element in a list."""
    code = "lst = [1, gather(a(), b()), 3]"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.7μs -> 8.79μs (45.0% faster)

def test_gather_as_keyword_argument(remover):
    """Test gather call as a keyword argument to another function."""
    code = "foo(bar=gather(a(), b()))"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.8μs -> 7.59μs (68.3% faster)

def test_gather_in_docstring(remover):
    """Test 'gather' in a docstring, not as a function call."""
    code = 'def foo():\n    """This function uses gather."""\n    pass'
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 11.9μs -> 10.1μs (17.6% faster)

def test_gather_as_attribute_of_other(remover):
    """Test gather as attribute of another object (not asyncio)."""
    code = "obj.gather(a(), b())"
    tree = ast.parse(code)
    # Should not match, as only 'asyncio.gather' or 'gather'
    codeflash_output = not remover._contains_asyncio_gather(tree) # 15.6μs -> 14.3μs (9.31% faster)

def test_gather_in_comment(remover):
    """Test 'gather' in a comment, not as a function call."""
    code = "def foo():\n    # gather(a(), b())\n    pass"
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 9.94μs -> 8.31μs (19.5% faster)

# ----------------------
# Large Scale Test Cases
# ----------------------

def test_large_number_of_functions_with_one_gather(remover):
    """Test a large module with many functions, only one has a gather call."""
    code = "\n".join([
        f"def func{i}():\n    return {i}" for i in range(500)
    ] + [
        "def gather_func():\n    return gather(a(), b())"
    ])
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 2.16ms -> 34.9μs (6104% faster)

def test_large_number_of_gather_calls(remover):
    """Test a module with many gather calls."""
    code = "\n".join([
        f"def func{i}():\n    return gather(a{i}(), b{i}())" for i in range(500)
    ])
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 1.82ms -> 36.2μs (4929% faster)

def test_large_number_of_asyncio_gather_calls(remover):
    """Test a module with many asyncio.gather calls."""
    code = "\n".join([
        f"def func{i}():\n    return asyncio.gather(a{i}(), b{i}())" for i in range(500)
    ])
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 1.82ms -> 36.0μs (4951% faster)

def test_large_code_with_no_gather(remover):
    """Test a large module with no gather calls."""
    code = "\n".join([
        f"def func{i}():\n    return a{i}() + b{i}()" for i in range(500)
    ])
    tree = ast.parse(code)
    codeflash_output = not remover._contains_asyncio_gather(tree) # 4.61ms -> 4.33ms (6.62% faster)

def test_large_mixed_code(remover):
    """Test a module with many functions, some with gather, some without."""
    code = "\n".join([
        f"def func{i}():\n    return gather(a{i}(), b{i}())" if i % 10 == 0 else f"def func{i}():\n    return a{i}()"
        for i in range(500)
    ])
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 1.82ms -> 87.3μs (1983% faster)
# 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 AsyncIOGatherRemover

# unit tests

@pytest.fixture
def remover():
    """Fixture for AsyncIOGatherRemover instance."""
    return AsyncIOGatherRemover()

# ------------------------
# Basic Test Cases
# ------------------------

def test_no_function_calls(remover):
    # Test with code that has no function calls
    code = "x = 1 + 2"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.1μs -> 10.9μs (11.0% faster)

def test_simple_asyncio_gather_call(remover):
    # Test with a direct asyncio.gather call
    code = "asyncio.gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 9.40μs -> 5.00μs (88.0% faster)

def test_simple_gather_call(remover):
    # Test with a bare gather call (not qualified)
    code = "gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 9.24μs -> 4.69μs (97.0% faster)

def test_other_asyncio_method(remover):
    # Test with another asyncio method, not gather
    code = "asyncio.sleep(1)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 11.9μs -> 11.0μs (8.67% faster)

def test_other_method_named_gather(remover):
    # Test with a method gather on another object
    code = "foo.gather(1, 2)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.5μs -> 11.7μs (6.40% faster)

# ------------------------
# Edge Test Cases
# ------------------------

def test_nested_gather_call(remover):
    # Test with a nested asyncio.gather call inside a function
    code = """
def foo():
    return asyncio.gather(a, b)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.1μs -> 7.28μs (79.8% faster)

def test_gather_in_lambda(remover):
    # Test with a gather call inside a lambda function
    code = "f = lambda: asyncio.gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 14.0μs -> 6.69μs (109% faster)

def test_gather_as_variable_name(remover):
    # Test with 'gather' used as a variable, not a function
    code = "gather = 42"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 9.26μs -> 7.41μs (24.9% faster)

def test_gather_as_function_argument(remover):
    # Test with 'gather' passed as an argument, not called
    code = "foo(gather)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 10.4μs -> 9.08μs (15.1% faster)

def test_gather_as_keyword_argument(remover):
    # Test with 'gather' as a keyword argument name
    code = "foo(gather=5)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 11.5μs -> 9.86μs (16.5% faster)

def test_gather_in_class_method(remover):
    # Test with asyncio.gather called inside a class method
    code = """
class C:
    def m(self):
        return asyncio.gather(a, b)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.7μs -> 8.82μs (78.4% faster)

def test_gather_in_async_function(remover):
    # Test with asyncio.gather in an async function
    code = """
async def foo():
    await asyncio.gather(a, b)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 14.0μs -> 7.78μs (79.8% faster)

def test_gather_in_comprehension(remover):
    # Test with gather used inside a list comprehension
    code = "[gather(x) for x in range(3)]"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 9.97μs -> 13.1μs (23.7% slower)

def test_gather_in_try_except(remover):
    # Test with asyncio.gather in a try-except block
    code = """
try:
    asyncio.gather(a, b)
except Exception:
    pass
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 12.3μs -> 10.5μs (17.3% faster)

def test_gather_in_decorator(remover):
    # Test with gather used as a decorator
    code = """
@gather
def foo(): pass
"""
    tree = ast.parse(code)
    # Should be False, because gather is not called, just used as a decorator
    codeflash_output = remover._contains_asyncio_gather(tree) # 11.7μs -> 10.3μs (13.5% faster)

def test_gather_as_attribute_of_other(remover):
    # Test with gather as an attribute of a non-asyncio object
    code = "obj.gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.1μs -> 12.1μs (8.81% faster)

def test_gather_as_function_definition(remover):
    # Test with a function named gather being defined
    code = """
def gather(x, y):
    return x + y
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 16.5μs -> 15.8μs (3.86% faster)

def test_asyncio_gather_with_alias(remover):
    # Test with asyncio imported as an alias
    code = """
import asyncio as aio
aio.gather(a, b)
"""
    tree = ast.parse(code)
    # Should be False, because the code only matches 'asyncio.gather', not 'aio.gather'
    codeflash_output = remover._contains_asyncio_gather(tree) # 15.1μs -> 13.7μs (10.4% faster)

def test_gather_with_different_casing(remover):
    # Test with gather called with different casing
    code = "asyncio.Gather(a, b)"
    tree = ast.parse(code)
    # Should be False, because 'Gather' != 'gather'
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.1μs -> 12.0μs (8.92% faster)

def test_gather_in_comment(remover):
    # Test with gather mentioned in a comment
    code = "# asyncio.gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 5.01μs -> 2.58μs (94.5% faster)

def test_gather_in_string(remover):
    # Test with gather mentioned in a string
    code = "x = 'asyncio.gather(a, b)'"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 8.98μs -> 7.28μs (23.3% faster)

def test_multiple_calls_some_gather(remover):
    # Test with multiple function calls, only some are gather
    code = """
foo(a)
bar(b)
asyncio.gather(a, b)
baz(c)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 13.5μs -> 9.62μs (40.1% faster)

def test_multiple_calls_none_gather(remover):
    # Test with multiple function calls, none are gather
    code = """
foo(a)
bar(b)
baz(c)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 18.4μs -> 17.5μs (5.68% faster)

def test_gather_in_nested_scopes(remover):
    # Test with gather call deeply nested in control flow
    code = """
def foo():
    if True:
        for i in range(2):
            if i == 1:
                asyncio.gather(a, b)
"""
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 26.5μs -> 11.7μs (126% faster)

# ------------------------
# Large Scale Test Cases
# ------------------------

def test_large_number_of_non_gather_calls(remover):
    # Test with a large number of function calls, none are gather
    code = "\n".join(f"foo{i}()" for i in range(500))
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 1.39ms -> 1.28ms (8.19% faster)

def test_large_number_of_gather_calls(remover):
    # Test with a large number of asyncio.gather calls
    code = "\n".join(f"asyncio.gather(a{i}, b{i})" for i in range(500))
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 313μs -> 32.6μs (864% faster)

def test_large_mixed_calls_with_one_gather(remover):
    # Test with many non-gather calls and one gather call at the end
    code = "\n".join(f"foo{i}()" for i in range(499)) + "\nasyncio.gather(a, b)"
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 840μs -> 32.6μs (2479% faster)

def test_large_mixed_calls_with_one_gather_at_start(remover):
    # Test with one gather call at the start and many non-gather calls
    code = "asyncio.gather(a, b)\n" + "\n".join(f"foo{i}()" for i in range(499))
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 313μs -> 1.28ms (75.5% slower)

def test_large_mixed_calls_with_bare_gather(remover):
    # Test with many non-gather calls and one bare gather call in the middle
    lines = [f"foo{i}()" for i in range(250)]
    lines.append("gather(a, b)")
    lines.extend(f"foo{i}()" for i in range(250, 500))
    code = "\n".join(lines)
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 574μs -> 654μs (12.1% slower)

def test_large_code_with_no_calls(remover):
    # Test with a large code block with only assignments (no function calls)
    code = "\n".join(f"x{i} = {i}" for i in range(1000))
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 2.79ms -> 2.63ms (6.30% faster)

def test_large_code_with_comments_and_gather(remover):
    # Test with a large code block, comments, and a single gather call
    lines = [f"# comment {i}" for i in range(500)]
    lines.append("asyncio.gather(a, b)")
    code = "\n".join(lines)
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 10.0μs -> 5.19μs (93.0% faster)

def test_large_code_with_strings_and_gather(remover):
    # Test with a large code block with gather in strings and one real gather call
    lines = [f"x = 'asyncio.gather({i}, {i+1})'" for i in range(500)]
    lines.append("asyncio.gather(a, b)")
    code = "\n".join(lines)
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 1.22ms -> 33.1μs (3601% faster)

def test_large_code_with_only_gather_in_strings(remover):
    # Test with a large code block with gather only in strings, never called
    code = "\n".join(f"x = 'asyncio.gather({i}, {i+1})'" for i in range(1000))
    tree = ast.parse(code)
    codeflash_output = remover._contains_asyncio_gather(tree) # 2.82ms -> 2.62ms (7.57% faster)
# 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-pr617-2025-09-02T17.32.23

Click to see suggested changes
Suggested change
for child_node in ast.walk(node):
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Attribute)
and isinstance(child_node.func.value, ast.Name)
and child_node.func.value.id == "asyncio"
and child_node.func.attr == "gather"
):
return True
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Name)
and child_node.func.id == "gather"
):
return True
stack = [node]
while stack:
child_node = stack.pop()
# Check for `asyncio.gather`
if isinstance(child_node, ast.Call):
func = child_node.func
if (
isinstance(func, ast.Attribute)
and isinstance(func.value, ast.Name)
and func.value.id == "asyncio"
and func.attr == "gather"
):
return True
# Check for direct `gather`
if isinstance(func, ast.Name) and func.id == "gather":
return True
stack.extend(ast.iter_child_nodes(child_node))

Comment on lines +54 to +70
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 36% (0.36x) speedup for AsyncIOGatherRemover.visit_FunctionDef in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 5.23 milliseconds → 3.84 milliseconds (best of 85 runs)

📝 Explanation and details

The optimization replaces ast.walk() with a manual stack-based traversal in the _contains_asyncio_gather() method. The key improvement comes from early termination - the optimized version can immediately return True when it finds an asyncio.gather call, while ast.walk() always traverses the entire AST regardless of early returns.

Key changes:

  • Stack-based traversal: Uses stack = [node] with stack.pop() and stack.extend(ast.iter_child_nodes(child_node))
  • Early exit optimization: When a gather call is found, the function returns immediately without processing remaining nodes
  • Reduced function call overhead: Eliminates the generator overhead from ast.walk()

Why it's faster:
The profiler shows the original code spent 81.7% of time in ast.walk() (54.8ms), while the optimized version spends only 82.4% of time in stack.extend() (32.5ms) - a significant reduction. More importantly, when gather calls are found early in the AST, the optimization provides dramatic speedups by avoiding unnecessary traversal.

Performance by test case:

  • Best for functions with gather calls: Tests with asyncio.gather or direct gather() calls show 50-200% speedups because early termination prevents full AST traversal
  • Modest gains for gather-free functions: 3-11% improvement from reduced function call overhead
  • Exceptional gains for large functions with gather: The large function test shows 3206% speedup, demonstrating how early termination scales with AST size

The optimization is particularly effective when the target pattern (gather calls) appears early in the AST structure, as it can skip processing hundreds or thousands of remaining nodes.

✅ Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 62 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 AsyncIOGatherRemover

# unit tests

# Helper function to parse source and extract the first FunctionDef node
def get_first_functiondef(source: str) -> ast.FunctionDef:
    """Parse Python source and return the first FunctionDef node."""
    tree = ast.parse(source)
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, ast.FunctionDef):
            return node
    raise ValueError("No FunctionDef found in source")

# -----------------
# 1. Basic Test Cases
# -----------------

def test_non_test_function_no_gather_retains_function():
    """Function not starting with 'test_' and no gather should be kept."""
    src = "def foo():\n    x = 1"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 16.6μs -> 17.2μs (3.50% slower)

def test_test_function_no_gather_retains_function():
    """Function starting with 'test_' but no gather should be kept."""
    src = "def test_example():\n    y = 2"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 26.9μs -> 25.1μs (7.26% faster)

def test_test_function_with_asyncio_gather_removes_function():
    """Function starting with 'test_' and containing asyncio.gather should be removed."""
    src = "def test_async():\n    await asyncio.gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 13.0μs -> 6.49μs (100% faster)

def test_test_function_with_gather_direct_call_removes_function():
    """Function starting with 'test_' and containing direct gather() call should be removed."""
    src = "def test_direct():\n    await gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 12.4μs -> 6.14μs (101% faster)

# -----------------
# 2. Edge Test Cases
# -----------------

def test_non_test_function_with_asyncio_gather_kept():
    """Non-test function with asyncio.gather should not be removed."""
    src = "def foo():\n    await asyncio.gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 18.2μs -> 18.5μs (1.52% slower)

def test_test_function_with_gather_in_string_kept():
    """'gather' in a string literal should not trigger removal."""
    src = "def test_str():\n    x = 'asyncio.gather(a, b)'"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 26.1μs -> 24.5μs (6.41% faster)

def test_test_function_with_gather_as_variable_kept():
    """'gather' used as a variable, not a call, should not trigger removal."""
    src = "def test_var():\n    gather = 5\n    x = gather + 2"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 41.0μs -> 39.5μs (3.75% faster)

def test_test_function_with_gather_call_on_other_object_kept():
    """gather called as method on another object should not trigger removal."""
    src = "def test_other():\n    obj.gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 31.3μs -> 29.8μs (5.11% faster)

def test_test_function_with_gather_in_comment_kept():
    """gather in a comment should not trigger removal."""
    src = "def test_comment():\n    # asyncio.gather(a, b)\n    pass"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 15.9μs -> 14.2μs (11.6% faster)

def test_test_function_with_nested_gather_call_removes_function():
    """gather call nested in an inner function should trigger removal."""
    src = (
        "def test_nested():\n"
        "    def inner():\n"
        "        asyncio.gather(a, b)\n"
        "    inner()"
    )
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 17.0μs -> 10.8μs (58.1% faster)

def test_test_function_with_gather_in_lambda_removes_function():
    """gather call inside a lambda triggers removal."""
    src = "def test_lambda():\n    f = lambda: asyncio.gather(a, b)\n    f()"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 18.6μs -> 10.6μs (75.0% faster)

def test_test_function_with_gather_as_keyword_arg_kept():
    """'gather' used as a keyword argument should not trigger removal."""
    src = "def test_kwarg():\n    foo(gather=5)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 30.0μs -> 28.8μs (4.07% faster)

def test_test_function_with_gather_imported_as_alias_removes_function():
    """Direct gather() call, even if imported as alias, triggers removal."""
    src = "def test_alias():\n    await gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 12.4μs -> 6.11μs (102% faster)

def test_test_function_with_gather_in_try_except_removes_function():
    """gather call inside try/except triggers removal."""
    src = (
        "def test_try():\n"
        "    try:\n"
        "        asyncio.gather(a, b)\n"
        "    except Exception:\n"
        "        pass"
    )
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 14.3μs -> 10.6μs (35.1% faster)

def test_test_function_with_gather_in_listcomp_removes_function():
    """gather call inside a list comprehension triggers removal."""
    src = "def test_listcomp():\n    [asyncio.gather(a, b) for a in range(2)]"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 12.8μs -> 13.1μs (2.30% slower)

def test_test_function_with_gather_in_decorator_kept():
    """gather used in a decorator name should not trigger removal."""
    src = "@gather\ndef test_decorated():\n    pass"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 19.9μs -> 18.4μs (7.99% faster)

def test_test_function_with_gather_in_attribute_but_not_asyncio_kept():
    """gather called as attribute of something other than 'asyncio' should not trigger removal."""
    src = "def test_attr():\n    foo.gather(a, b)"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 31.4μs -> 29.7μs (5.57% faster)

def test_test_function_with_gather_in_multiple_places_removes_function():
    """Multiple gather calls in a test function should trigger removal."""
    src = (
        "def test_multi():\n"
        "    asyncio.gather(a, b)\n"
        "    gather(a, b)\n"
        "    x = 1"
    )
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 13.6μs -> 9.06μs (50.2% faster)

def test_test_function_with_gather_as_class_method_removes_function():
    """gather call inside a test function as a class method triggers removal."""
    src = (
        "def test_class():\n"
        "    class C:\n"
        "        def m(self):\n"
        "            asyncio.gather(a, b)\n"
        "    C().m()"
    )
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 20.8μs -> 14.8μs (40.6% faster)

# -----------------
# 3. Large Scale Test Cases
# -----------------

def test_large_function_with_many_statements_and_one_gather_removes_function():
    """Large test function with many statements and one gather call should be removed."""
    body = "\n    ".join([f"x{i} = {i}" for i in range(500)])
    src = f"def test_large():\n    {body}\n    asyncio.gather(a, b)\n    y = 3"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 1.24ms -> 37.4μs (3206% faster)

def test_large_function_with_many_statements_no_gather_kept():
    """Large test function with many statements, no gather call, should be kept."""
    body = "\n    ".join([f"x{i} = {i}" for i in range(500)])
    src = f"def test_large_nogather():\n    {body}\n    y = 3"
    node = get_first_functiondef(src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_FunctionDef(node); result = codeflash_output # 3.60ms -> 3.46ms (3.88% faster)

def test_many_functions_some_with_gather():
    """Apply remover to a module with many functions, some test_ with gather, some not."""
    src = (
        "def test_a():\n    asyncio.gather(a, b)\n"
        "def test_b():\n    pass\n"
        "def foo():\n    asyncio.gather(a, b)\n"
        "def test_c():\n    gather(a, b)\n"
        "def bar():\n    pass\n"
        "def test_d():\n    pass\n"
    )
    tree = ast.parse(src)
    remover = AsyncIOGatherRemover()
    new_tree = remover.visit(tree)
    # Collect function names in the transformed tree
    kept_names = [n.name for n in new_tree.body if isinstance(n, ast.FunctionDef)]

def test_large_number_of_functions_performance():
    """Performance: 1000 functions, only a few test_ with gather should be removed."""
    src_lines = []
    # 995 non-test functions
    for i in range(995):
        src_lines.append(f"def func_{i}():\n    pass")
    # 5 test functions with gather
    for i in range(5):
        src_lines.append(f"def test_remove_{i}():\n    asyncio.gather(a, b)")
    src = "\n".join(src_lines)
    tree = ast.parse(src)
    remover = AsyncIOGatherRemover()
    new_tree = remover.visit(tree)
    # Only the 995 non-test functions should remain
    kept_names = [n.name for n in new_tree.body if isinstance(n, ast.FunctionDef)]
# 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 AsyncIOGatherRemover

# unit tests

def parse_and_get_funcdef_nodes(source_code):
    """
    Helper to parse source code and return all ast.FunctionDef nodes.
    """
    tree = ast.parse(source_code)
    return [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)]

def run_remover_on_funcdef(node):
    """
    Helper to run AsyncIOGatherRemover on a single FunctionDef node.
    """
    remover = AsyncIOGatherRemover()
    return remover.visit_FunctionDef(node)

# 1. Basic Test Cases

def test_non_test_function_with_asyncio_gather_is_preserved():
    # Function does not start with 'test_' but contains asyncio.gather
    code = """
def foo():
    x = asyncio.gather(a(), b())
    return x
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_without_asyncio_gather_is_preserved():
    # Function starts with 'test_' but does NOT contain asyncio.gather
    code = """
def test_bar():
    x = await something()
    return x
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_asyncio_gather_is_removed():
    # Function starts with 'test_' and contains asyncio.gather
    code = """
def test_baz():
    await asyncio.gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_gather_imported_directly_is_removed():
    # Function starts with 'test_' and uses 'gather' directly (imported)
    code = """
def test_direct_gather():
    await gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_non_test_function_with_gather_directly_is_preserved():
    # Function does not start with 'test_' and uses 'gather' directly
    code = """
def foo_gather():
    await gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

# 2. Edge Test Cases

def test_test_function_with_asyncio_gather_in_nested_function_is_removed():
    # 'test_' function with gather in a nested function
    code = """
def test_nested():
    def inner():
        return asyncio.gather(a(), b())
    return inner()
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_gather_as_variable_not_removed():
    # 'test_' function where 'gather' is a variable, not a function call
    code = """
def test_var_gather():
    gather = 42
    return gather
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_similar_name_not_removed():
    # 'test_' function with 'gather' in a variable name, not a call
    code = """
def test_similar():
    gatherer = "something"
    return gatherer
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_asyncio_gather_as_attribute_of_other_object_not_removed():
    # 'test_' function with gather as attribute of something other than asyncio
    code = """
def test_other_gather():
    obj = SomeClass()
    await obj.gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_asyncio_gather_in_comment_not_removed():
    # 'test_' function with 'asyncio.gather' only in a comment
    code = """
def test_comment():
    # await asyncio.gather(a(), b())
    return 1
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_asyncio_gather_in_string_not_removed():
    # 'test_' function with 'asyncio.gather' only in a string
    code = '''
def test_string():
    s = "asyncio.gather(a(), b())"
    return s
'''
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_gather_called_on_different_object_not_removed():
    # 'test_' function with gather called on a different object (not asyncio)
    code = """
def test_other_obj_gather():
    await something.gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_function_with_name_test_but_not_at_start_is_preserved():
    # Function with 'test' in the name but not at the start
    code = """
def not_a_test_function():
    await asyncio.gather(a(), b())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_multiple_asyncio_gather_calls_is_removed():
    # 'test_' function with multiple asyncio.gather calls
    code = """
def test_multi_gather():
    await asyncio.gather(a(), b())
    await asyncio.gather(c(), d())
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_test_function_with_gather_as_keyword_argument_not_removed():
    # 'test_' function where 'gather' is a keyword argument, not a call
    code = """
def test_gather_kwarg():
    foo(gather=asyncio.gather)
"""
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

# 3. Large Scale Test Cases


def test_large_function_body_with_gather_only_once():
    # Large function body, gather appears only once
    body_lines = ["    x = %d" % i for i in range(400)]
    body_lines.append("    await asyncio.gather(a(), b())")
    body_lines.extend(["    y = %d" % i for i in range(400, 800)])
    code = "def test_large():\n" + "\n".join(body_lines)
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_large_function_body_no_gather():
    # Large function body, but no gather
    body_lines = ["    x = %d" % i for i in range(900)]
    code = "def test_large_no_gather():\n" + "\n".join(body_lines)
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

def test_large_number_of_nested_functions():
    # Many nested functions, only outermost is 'test_' with gather in deepest
    code = "def test_outer():\n"
    indent = "    "
    for i in range(50):
        code += f"{indent}def inner_{i}():\n"
        indent += "    "
    code += f"{indent}return asyncio.gather(a(), b())\n"
    func_node = parse_and_get_funcdef_nodes(code)[0]
    result = run_remover_on_funcdef(func_node)

To test or edit this optimization locally git merge codeflash/optimize-pr617-2025-09-02T17.39.22

Click to see suggested changes
Suggested change
for child_node in ast.walk(node):
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Attribute)
and isinstance(child_node.func.value, ast.Name)
and child_node.func.value.id == "asyncio"
and child_node.func.attr == "gather"
):
return True
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Name)
and child_node.func.id == "gather"
):
return True
stack = [node]
while stack:
child_node = stack.pop()
if isinstance(child_node, ast.Call):
func = child_node.func
if (
isinstance(func, ast.Attribute)
and isinstance(func.value, ast.Name)
and func.value.id == "asyncio"
and func.attr == "gather"
):
return True
if isinstance(func, ast.Name) and func.id == "gather":
return True
stack.extend(ast.iter_child_nodes(child_node))

Comment on lines +54 to +70
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚡️Codeflash found 25% (0.25x) speedup for AsyncIOGatherRemover.visit_AsyncFunctionDef in codeflash/code_utils/instrument_existing_tests.py

⏱️ Runtime : 781 microseconds → 627 microseconds (best of 71 runs)

📝 Explanation and details

The optimization replaces ast.walk() with a custom iterative depth-first search using a stack, which provides early termination when asyncio.gather calls are found.

Key optimizations:

  1. Short-circuit traversal: The original code uses ast.walk(node) which always traverses the entire AST subtree, even after finding a match. The optimized version uses an explicit stack and immediately returns True upon finding the first gather call, avoiding unnecessary node visits.

  2. Iterative over recursive: Replaces the generator-based ast.walk() with a simple while loop and stack operations (stack.pop(), stack.extend()), reducing function call overhead.

  3. Streamlined control flow: The nested if conditions are restructured to check ast.Call once, then branch to handle the two gather patterns (attribute vs direct call), eliminating redundant isinstance checks.

Performance impact based on test results:

  • Most effective on test functions that contain gather calls (103-136% speedup) because early termination kicks in immediately
  • Moderate gains on functions without gather calls (6-10% speedup) due to reduced traversal overhead
  • Minimal impact on non-test functions (2-4% speedup) since the name check happens first

The 24% overall speedup comes from avoiding the continuation of AST traversal after matches are found, which is particularly beneficial when gather calls appear early in the function body or in nested structures.

✅ Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 74 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 AsyncIOGatherRemover

# unit tests

def make_async_func(name, body_src):
    """Helper to create an AsyncFunctionDef node from code string."""
    src = f"async def {name}():\n"
    # Indent body
    body_src = "\n".join("    " + line if line.strip() else line for line in body_src.splitlines())
    src += body_src
    mod = ast.parse(src)
    # Return the AsyncFunctionDef node
    for node in mod.body:
        if isinstance(node, ast.AsyncFunctionDef):
            return node
    raise ValueError("No AsyncFunctionDef found")

# ------------------ Basic Test Cases ------------------

def test_basic_removes_test_with_asyncio_gather():
    # Should remove test_ functions containing asyncio.gather
    node = make_async_func("test_example", "await asyncio.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 15.0μs -> 7.40μs (103% faster)

def test_basic_keeps_non_test_with_asyncio_gather():
    # Should NOT remove non-test_ functions even if they have asyncio.gather
    node = make_async_func("run_example", "await asyncio.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 19.2μs -> 18.8μs (2.46% faster)

def test_basic_keeps_test_without_asyncio_gather():
    # Should NOT remove test_ functions if they don't have asyncio.gather
    node = make_async_func("test_no_gather", "await something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 30.9μs -> 28.3μs (9.23% faster)

def test_basic_keeps_non_test_without_asyncio_gather():
    # Should NOT remove non-test_ functions without asyncio.gather
    node = make_async_func("run_no_gather", "await something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 16.7μs -> 16.3μs (2.77% faster)

def test_basic_removes_test_with_gather_imported():
    # Should remove test_ functions using imported gather
    node = make_async_func("test_imported_gather", "await gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 13.9μs -> 6.40μs (117% faster)

# ------------------ Edge Test Cases ------------------

def test_edge_test_name_case_sensitive():
    # Should only remove if name starts with lowercase 'test_'
    node = make_async_func("Test_example", "await asyncio.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 18.5μs -> 17.9μs (3.31% faster)

def test_edge_gather_as_variable():
    # Should NOT remove if 'gather' is used as a variable, not a function
    node = make_async_func("test_var_gather", "gather = 5\nawait something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 42.6μs -> 39.2μs (8.74% faster)

def test_edge_gather_in_comment():
    # Should NOT remove if 'asyncio.gather' appears only in a comment
    node = make_async_func("test_comment_gather", "# await asyncio.gather(a, b)\nawait something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 30.4μs -> 27.6μs (10.2% faster)

def test_edge_gather_in_string():
    # Should NOT remove if 'asyncio.gather' appears only in a string
    node = make_async_func("test_string_gather", "x = 'await asyncio.gather(a, b)'\nawait something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 42.0μs -> 38.4μs (9.34% faster)

def test_edge_nested_gather_call():
    # Should remove if asyncio.gather is nested deep in the function
    node = make_async_func("test_nested_gather", "if True:\n    await asyncio.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 16.0μs -> 7.49μs (114% faster)

def test_edge_multiple_gather_calls():
    # Should remove if multiple gather calls are present
    node = make_async_func("test_multi_gather", "await asyncio.gather(a, b)\nawait gather(c, d)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 14.7μs -> 6.22μs (136% faster)

def test_edge_gather_as_method_of_other_object():
    # Should NOT remove if gather is called as method of something other than asyncio
    node = make_async_func("test_other_gather", "await foo.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 33.4μs -> 30.6μs (8.92% faster)

def test_edge_gather_in_lambda():
    # Should remove if asyncio.gather is called inside a lambda
    node = make_async_func("test_lambda_gather", "x = lambda: asyncio.gather(a, b)\nawait something_else(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 18.8μs -> 13.6μs (38.7% faster)

def test_edge_gather_in_inner_function():
    # Should remove if asyncio.gather is called in an inner function
    src = (
        "def inner():\n"
        "    return asyncio.gather(a, b)\n"
        "await something_else(a, b)"
    )
    node = make_async_func("test_inner_func_gather", src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 17.6μs -> 13.8μs (26.9% faster)

def test_edge_gather_in_try_except():
    # Should remove if asyncio.gather is called inside try/except
    src = (
        "try:\n"
        "    await asyncio.gather(a, b)\n"
        "except Exception:\n"
        "    pass"
    )
    node = make_async_func("test_try_gather", src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 17.4μs -> 10.9μs (58.8% faster)

def test_edge_gather_in_async_for():
    # Should remove if asyncio.gather is called inside async for
    src = (
        "async for x in y:\n"
        "    await asyncio.gather(a, b)"
    )
    node = make_async_func("test_async_for_gather", src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 17.6μs -> 7.91μs (123% faster)

def test_edge_gather_in_async_with():
    # Should remove if asyncio.gather is called inside async with
    src = (
        "async with foo:\n"
        "    await asyncio.gather(a, b)"
    )
    node = make_async_func("test_async_with_gather", src)
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 17.1μs -> 7.59μs (125% faster)

def test_edge_gather_with_kwargs():
    # Should remove if asyncio.gather is called with keyword arguments
    node = make_async_func("test_gather_kwargs", "await asyncio.gather(a, b, return_exceptions=True)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 13.3μs -> 6.08μs (119% faster)

def test_edge_gather_with_star_args():
    # Should remove if asyncio.gather is called with *args
    node = make_async_func("test_gather_star", "await asyncio.gather(*tasks)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 13.1μs -> 6.16μs (113% faster)

def test_edge_gather_with_alias():
    # Should NOT remove if asyncio is imported as something else and used
    node = make_async_func("test_alias_gather", "import asyncio as aio\nawait aio.gather(a, b)")
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 37.9μs -> 35.3μs (7.62% faster)

# ------------------ Large Scale Test Cases ------------------







#------------------------------------------------
from __future__ import annotations

import ast

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.instrument_existing_tests import AsyncIOGatherRemover

# unit tests

# Helper to parse code and get AsyncFunctionDef node(s)
def get_async_func_nodes(code: str):
    """Parse code and return all AsyncFunctionDef nodes."""
    tree = ast.parse(code)
    return [n for n in ast.walk(tree) if isinstance(n, ast.AsyncFunctionDef)]

# Basic Test Cases

def test_basic_removal_with_asyncio_gather():
    # Should remove test function with asyncio.gather
    code = """
import asyncio
async def test_func():
    await asyncio.gather(foo(), bar())
    return 1
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 12.5μs -> 7.37μs (69.6% faster)

def test_basic_removal_with_gather_direct():
    # Should remove test function with direct gather call
    code = """
async def test_func():
    await gather(foo(), bar())
    return 2
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 12.3μs -> 7.04μs (75.1% faster)

def test_basic_no_removal_non_test_func():
    # Should NOT remove non-test function even with asyncio.gather
    code = """
import asyncio
async def not_a_test():
    await asyncio.gather(foo(), bar())
    return 3
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 26.0μs -> 25.9μs (0.545% faster)

def test_basic_no_removal_test_func_no_gather():
    # Should NOT remove test function without asyncio.gather
    code = """
async def test_func():
    await foo()
    return 4
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 29.1μs -> 27.3μs (6.45% faster)

# Edge Test Cases

def test_edge_removal_gather_in_nested_function():
    # Should remove if gather is in a nested function inside test_
    code = """
async def test_func():
    async def inner():
        await asyncio.gather(foo())
    await inner()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 16.1μs -> 10.5μs (52.5% faster)

def test_edge_removal_gather_in_lambda():
    # Should remove if gather is called inside a lambda in test_
    code = """
async def test_func():
    x = lambda: asyncio.gather(foo())
    await foo()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 15.4μs -> 9.83μs (56.6% faster)

def test_edge_removal_gather_in_comprehension():
    # Should remove if gather is used in a list comprehension
    code = """
async def test_func():
    results = [asyncio.gather(foo()) for _ in range(2)]
    await foo()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 14.1μs -> 15.4μs (8.54% slower)

def test_edge_removal_gather_in_decorator():
    # Should remove if gather is used in a decorator
    code = """
def dec(func):
    async def wrapper(*a, **k):
        await asyncio.gather(foo())
        return await func(*a, **k)
    return wrapper

@dec
async def test_func():
    await foo()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 25.4μs -> 23.2μs (9.41% faster)

def test_edge_removal_gather_as_variable_name():
    # Should NOT remove if 'gather' is used as variable name
    code = """
async def test_func():
    gather = 42
    await foo()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 33.4μs -> 31.4μs (6.31% faster)

def test_edge_removal_gather_as_attribute_not_asyncio():
    # Should NOT remove if gather is called as attribute of something else
    code = """
async def test_func():
    await mylib.gather(foo())
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 29.7μs -> 28.6μs (3.85% faster)

def test_edge_removal_gather_in_comment_only():
    # Should NOT remove if gather is only mentioned in a comment
    code = """
async def test_func():
    # await asyncio.gather(foo())
    await foo()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 21.5μs -> 20.2μs (6.41% faster)

def test_edge_removal_gather_in_docstring_only():
    # Should NOT remove if gather is only mentioned in a docstring
    code = '''
async def test_func():
    """This function uses asyncio.gather in the future."""
    await foo()
'''
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 28.2μs -> 26.6μs (5.91% faster)

def test_edge_removal_test_func_name_variants():
    # Should remove if name starts with test_ exactly
    code = """
async def test_foo():
    await asyncio.gather(foo())
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 10.6μs -> 5.32μs (99.3% faster)

    # Should NOT remove if name does not start with test_
    code2 = """
async def testfoo():
    await asyncio.gather(foo())
"""
    node2 = get_async_func_nodes(code2)[0]
    remover2 = AsyncIOGatherRemover()
    codeflash_output = remover2.visit_AsyncFunctionDef(node2); result2 = codeflash_output # 16.3μs -> 15.6μs (4.17% faster)


def test_edge_removal_gather_called_with_no_args():
    # Should remove even if gather is called with no arguments
    code = """
async def test_func():
    await asyncio.gather()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 11.1μs -> 5.79μs (91.7% faster)

def test_edge_removal_gather_called_as_keyword():
    # Should remove if gather is called with keywords only
    code = """
async def test_func():
    await asyncio.gather(a=foo())
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 10.7μs -> 5.55μs (93.1% faster)

# Large Scale Test Cases


def test_large_scale_nested_gather_calls():
    # Should remove deeply nested gather calls in test_ functions
    code = """
async def test_func():
    async def a():
        async def b():
            await asyncio.gather(foo())
        await b()
    await a()
"""
    node = get_async_func_nodes(code)[0]
    remover = AsyncIOGatherRemover()
    codeflash_output = remover.visit_AsyncFunctionDef(node); result = codeflash_output # 22.8μs -> 15.6μs (46.3% faster)

To test or edit this optimization locally git merge codeflash/optimize-pr617-2025-09-02T17.47.10

Click to see suggested changes
Suggested change
for child_node in ast.walk(node):
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Attribute)
and isinstance(child_node.func.value, ast.Name)
and child_node.func.value.id == "asyncio"
and child_node.func.attr == "gather"
):
return True
if (
isinstance(child_node, ast.Call)
and isinstance(child_node.func, ast.Name)
and child_node.func.id == "gather"
):
return True
stack = [node]
while stack:
current = stack.pop()
if isinstance(current, ast.Call):
if isinstance(current.func, ast.Attribute):
if (
isinstance(current.func.value, ast.Name)
and current.func.value.id == "asyncio"
and current.func.attr == "gather"
):
return True
elif isinstance(current.func, ast.Name) and current.func.id == "gather":
return True
stack.extend(ast.iter_child_nodes(current))

return False

def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef | None:
if node.name.startswith("test_") and self._contains_asyncio_gather(node):
return None
return self.generic_visit(node)

def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef | None:
if node.name.startswith("test_") and self._contains_asyncio_gather(node):
return None
return self.generic_visit(node)


class InjectPerfOnly(ast.NodeTransformer):
def __init__(
self,
Expand All @@ -56,6 +89,8 @@ def __init__(
test_framework: str,
call_positions: list[CodePosition],
mode: TestingMode = TestingMode.BEHAVIOR,
*,
is_async: bool = False,
) -> None:
self.mode: TestingMode = mode
self.function_object = function
Expand All @@ -64,13 +99,16 @@ def __init__(
self.module_path = module_path
self.test_framework = test_framework
self.call_positions = call_positions
self.is_async = is_async
if len(function.parents) == 1 and function.parents[0].type == "ClassDef":
self.class_name = function.top_level_parent_name

def find_and_update_line_node(
self, test_node: ast.stmt, node_name: str, index: str, test_class_name: str | None = None
) -> Iterable[ast.stmt] | None:
call_node = None
await_node = None

for node in ast.walk(test_node):
if isinstance(node, ast.Call) and node_in_call_position(node, self.call_positions):
call_node = node
Expand Down Expand Up @@ -120,6 +158,64 @@ def find_and_update_line_node(
node.keywords = call_node.keywords
break

# Check for awaited function calls
elif (
isinstance(node, ast.Await)
and isinstance(node.value, ast.Call)
and node_in_call_position(node.value, self.call_positions)
):
call_node = node.value
await_node = node
if isinstance(call_node.func, ast.Name):
function_name = call_node.func.id
call_node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
call_node.args = [
ast.Name(id=function_name, ctx=ast.Load()),
ast.Constant(value=self.module_path),
ast.Constant(value=test_class_name or None),
ast.Constant(value=node_name),
ast.Constant(value=self.function_object.qualified_name),
ast.Constant(value=index),
ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
*(
[ast.Name(id="codeflash_cur", ctx=ast.Load()), ast.Name(id="codeflash_con", ctx=ast.Load())]
if self.mode == TestingMode.BEHAVIOR
else []
),
*call_node.args,
]
call_node.keywords = call_node.keywords
# Keep the await wrapper around the modified call
await_node.value = call_node
break
if isinstance(call_node.func, ast.Attribute):
function_to_test = call_node.func.attr
if function_to_test == self.function_object.function_name:
function_name = ast.unparse(call_node.func)
call_node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
call_node.args = [
ast.Name(id=function_name, ctx=ast.Load()),
ast.Constant(value=self.module_path),
ast.Constant(value=test_class_name or None),
ast.Constant(value=node_name),
ast.Constant(value=self.function_object.qualified_name),
ast.Constant(value=index),
ast.Name(id="codeflash_loop_index", ctx=ast.Load()),
*(
[
ast.Name(id="codeflash_cur", ctx=ast.Load()),
ast.Name(id="codeflash_con", ctx=ast.Load()),
]
if self.mode == TestingMode.BEHAVIOR
else []
),
*call_node.args,
]
call_node.keywords = call_node.keywords
# Keep the await wrapper around the modified call
await_node.value = call_node
break

if call_node is None:
return None
return [test_node]
Expand All @@ -129,9 +225,34 @@ def visit_ClassDef(self, node: ast.ClassDef) -> ast.ClassDef:
for inner_node in ast.walk(node):
if isinstance(inner_node, ast.FunctionDef):
self.visit_FunctionDef(inner_node, node.name)
elif isinstance(inner_node, ast.AsyncFunctionDef):
self.visit_AsyncFunctionDef(inner_node, node.name)

return node

def visit_AsyncFunctionDef(
self, node: ast.AsyncFunctionDef, test_class_name: str | None = None
) -> ast.AsyncFunctionDef:
sync_node = ast.FunctionDef(
name=node.name,
args=node.args,
body=node.body,
decorator_list=node.decorator_list,
returns=node.returns,
lineno=node.lineno,
col_offset=node.col_offset if hasattr(node, "col_offset") else 0,
)
processed_sync = self.visit_FunctionDef(sync_node, test_class_name)
return ast.AsyncFunctionDef(
name=processed_sync.name,
args=processed_sync.args,
body=processed_sync.body,
decorator_list=processed_sync.decorator_list,
returns=processed_sync.returns,
lineno=processed_sync.lineno,
col_offset=processed_sync.col_offset if hasattr(processed_sync, "col_offset") else 0,
)

def visit_FunctionDef(self, node: ast.FunctionDef, test_class_name: str | None = None) -> ast.FunctionDef:
if node.name.startswith("test_"):
did_update = False
Expand Down Expand Up @@ -310,6 +431,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
file_path=self.function.file_path,
starting_line=self.function.starting_line,
ending_line=self.function.ending_line,
is_async=self.function.is_async,
)
else:
self.imported_as = FunctionToOptimize(
Expand All @@ -318,6 +440,7 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
file_path=self.function.file_path,
starting_line=self.function.starting_line,
ending_line=self.function.ending_line,
is_async=self.function.is_async,
)


Expand All @@ -342,7 +465,16 @@ def inject_profiling_into_existing_test(
import_visitor.visit(tree)
func = import_visitor.imported_as

tree = InjectPerfOnly(func, test_module_path, test_framework, call_positions, mode=mode).visit(tree)
is_async = function_to_optimize.is_async
logger.debug(f"Using async status from discovery phase for {function_to_optimize.function_name}: {is_async}")

if is_async:
asyncio_gather_remover = AsyncIOGatherRemover()
tree = asyncio_gather_remover.visit(tree)

tree = InjectPerfOnly(func, test_module_path, test_framework, call_positions, mode=mode, is_async=is_async).visit(
tree
)
new_imports = [
ast.Import(names=[ast.alias(name="time")]),
ast.Import(names=[ast.alias(name="gc")]),
Expand All @@ -354,11 +486,15 @@ 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]
if is_async:
new_imports.append(ast.Import(names=[ast.alias(name="inspect")]))
tree.body = [*new_imports, create_wrapper_function(mode, is_async=is_async), *tree.body]
return True, isort.code(ast.unparse(tree), float_to_top=True)


def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.FunctionDef:
def create_wrapper_function(
mode: TestingMode = TestingMode.BEHAVIOR, *, is_async: bool = False
) -> ast.FunctionDef | ast.AsyncFunctionDef:
lineno = 1
wrapper_body: list[ast.stmt] = [
ast.Assign(
Expand Down Expand Up @@ -534,14 +670,70 @@ def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.Fun
),
lineno=lineno + 11,
),
ast.Assign(
targets=[ast.Name(id="return_value", ctx=ast.Store())],
value=ast.Call(
func=ast.Name(id="wrapped", ctx=ast.Load()),
args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
),
lineno=lineno + 12,
# For async wrappers
# Call the wrapped function first, then check if result is awaitable before awaiting.
# This handles mixed scenarios where async tests might call both sync and async functions.
*(
[
ast.Assign(
targets=[ast.Name(id="ret", ctx=ast.Store())],
value=ast.Call(
func=ast.Name(id="wrapped", ctx=ast.Load()),
args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
),
lineno=lineno + 12,
),
ast.If(
test=ast.Call(
func=ast.Attribute(
value=ast.Name(id="inspect", ctx=ast.Load()), attr="isawaitable", ctx=ast.Load()
),
args=[ast.Name(id="ret", ctx=ast.Load())],
keywords=[],
),
body=[
ast.Assign(
targets=[ast.Name(id="counter", ctx=ast.Store())],
value=ast.Call(
func=ast.Attribute(
value=ast.Name(id="time", ctx=ast.Load()),
attr="perf_counter_ns",
ctx=ast.Load(),
),
args=[],
keywords=[],
),
lineno=lineno + 14,
),
ast.Assign(
targets=[ast.Name(id="return_value", ctx=ast.Store())],
value=ast.Await(value=ast.Name(id="ret", ctx=ast.Load())),
lineno=lineno + 15,
),
],
orelse=[
ast.Assign(
targets=[ast.Name(id="return_value", ctx=ast.Store())],
value=ast.Name(id="ret", ctx=ast.Load()),
lineno=lineno + 16,
)
],
lineno=lineno + 13,
),
]
if is_async
else [
ast.Assign(
targets=[ast.Name(id="return_value", ctx=ast.Store())],
value=ast.Call(
func=ast.Name(id="wrapped", ctx=ast.Load()),
args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
keywords=[ast.keyword(arg=None, value=ast.Name(id="kwargs", ctx=ast.Load()))],
),
lineno=lineno + 12,
)
]
),
ast.Assign(
targets=[ast.Name(id="codeflash_duration", ctx=ast.Store())],
Expand Down Expand Up @@ -703,7 +895,7 @@ def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.Fun
),
ast.Return(value=ast.Name(id="return_value", ctx=ast.Load()), lineno=lineno + 19),
]
return ast.FunctionDef(
func_def = ast.FunctionDef(
name="codeflash_wrap",
args=ast.arguments(
args=[
Expand All @@ -729,3 +921,13 @@ def create_wrapper_function(mode: TestingMode = TestingMode.BEHAVIOR) -> ast.Fun
decorator_list=[],
returns=None,
)
if is_async:
return ast.AsyncFunctionDef(
name="codeflash_wrap",
args=func_def.args,
body=func_def.body,
lineno=func_def.lineno,
decorator_list=func_def.decorator_list,
returns=func_def.returns,
)
return func_def
12 changes: 9 additions & 3 deletions codeflash/code_utils/static_analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,13 +128,19 @@ def get_first_top_level_object_def_ast(

def get_first_top_level_function_or_method_ast(
function_name: str, parents: list[FunctionParent], node: ast.AST
) -> ast.FunctionDef | None:
) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
if not parents:
return get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node)
result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, node)
if result is not None:
return result
return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, node)
if parents[0].type == "ClassDef" and (
class_node := get_first_top_level_object_def_ast(parents[0].name, ast.ClassDef, node)
):
return get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
result = get_first_top_level_object_def_ast(function_name, ast.FunctionDef, class_node)
if result is not None:
return result
return get_first_top_level_object_def_ast(function_name, ast.AsyncFunctionDef, class_node)
return None


Expand Down
Loading
Loading