⚡️ Speed up function find_codeflash_output_assignments by 21%

[codeflash-ai]

📄 21% (0.21x) speedup for find_codeflash_output_assignments in codeflash/code_utils/edit_generated_tests.py

⏱️ Runtime : 41.9 milliseconds → 34.7 milliseconds (best of 145 runs)

📝 Explanation and details

Here’s an optimized version of your code for speed, given the profiler shows almost all the time is inside ast.parse() and especially in visitor.visit(tree). The CfoVisitor implementation is not shown, but the code will optimize how it's used.

Changes/Optimizations.

• Eliminate source retention: Assuming CfoVisitor needs only line mappings, we avoid passing the entire source_code string if possible—if you do not control CfoVisitor, ignore this.
• Optimize AST walk: If CfoVisitor is a standard ast.NodeVisitor, using ast.walk() or writing an iterative visitor (instead of recursive) is sometimes faster. But if its internal logic is complex, you’d likely get the best speed by using a C-optimized AST walker, but Python's ast isn’t pluggable that way.
• Reuse AST parse trees if parsing the same file repeatedly.
• Disable docstring processing if not needed: ast.parse(..., type_comments=True) is not default; so we can’t skip more work there.
• ast.parse is already fast, but we can speed up repeated parses by using ast.parse with a support for possibly pre-compiled ASTs (not possible here since you always receive new source_code).

With the given constraints, the code is likely already close to optimal for pure Python. The only micro-optimization realistically available in this snippet is to.

• Avoid tracking line numbers if not needed,
• Remove the source_code string from being retained in the visitor if not required (saving memory, especially for huge inputs).

Most critically: If you control CfoVisitor, re-implementing its logic in Cython or with a C extension would vastly speed up the slowest portion (visit).

---

Assuming default use (no changes to CfoVisitor), here’s a slightly optimized code, including a minor AST traversal microspeedup.

If you control the implementation of CfoVisitor, remove storing source_code as an attribute if not needed, which can yield memory and minor speed improvements.

Further Optimization: Custom AST Walker

If CfoVisitor is simple, you may get a significant speed boost by rolling your own iterative AST traversal using ast.walk(tree), avoiding double-dispatch of visit_* methods (especially if only one node type is relevant).

---

Summary:

• The code provided is close to optimal unless you can optimize the logic in CfoVisitor.
• Use local variables for method lookups.
• Don't hold onto source code in memory unless needed.
• If you can, move your critical logic into a C or Cython extension for a big speed up.
• If only certain AST node types matter, iterate via ast.walk() and filter node types directly—reducing Python call overhead.

If you share the implementation of CfoVisitor, I can reimplement the whole logic to be even faster! But with your constraints, these are all possible micro-optimizations.

Let me know if you want a version that bakes in CfoVisitor logic or if you can share that code.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 89 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.edit_generated_tests import \
    find_codeflash_output_assignments

# unit tests

# ------------------ BASIC TEST CASES ------------------

def test_single_assignment():
    # Test a single assignment to output in a simple function
    code = """
def foo():
    output = 5
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_multiple_assignments():
    # Test multiple assignments to output in a function
    code = """
def foo():
    output = 1
    output = 2
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_if_else():
    # Test assignments to output in different branches
    code = """
def foo(x):
    if x > 0:
        output = 1
    else:
        output = -1
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_loop():
    # Test assignment inside a loop
    code = """
def foo(xs):
    for x in xs:
        output = x
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_with_type_annotation():
    # Test assignment with type annotation
    code = """
def foo():
    output: int = 42
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_augmented_assignment():
    # Test augmented assignment
    code = """
def foo():
    output = 0
    output += 5
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_tuple_assignment():
    # Test tuple assignment where output is one of the targets
    code = """
def foo():
    a, output, b = 1, 2, 3
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_to_other_variable():
    # Test that assignments to variables other than output are ignored
    code = """
def foo():
    result = 1
    return result
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_nested_function():
    # Test that assignment in nested function is not counted
    code = """
def foo():
    output = 1
    def bar():
        output = 2
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_nested_class():
    # Test that assignment in nested class is not counted
    code = """
def foo():
    output = 1
    class Bar:
        def baz(self):
            output = 2
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_lambda():
    # Test that assignment in lambda is not counted (shouldn't be possible, but test for coverage)
    code = """
def foo():
    output = 1
    f = lambda: (output := 2)
    return output
"""
    # Assignment expressions in lambdas to output should not be counted by our visitor
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_in_async_function():
    # Test function with async def
    code = """
async def foo():
    output = 1
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_within_with_statement():
    # Test assignment inside a with statement
    code = """
def foo():
    with open('file.txt') as f:
        output = f.read()
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_assignment_within_try_except():
    # Test assignment inside try/except/finally
    code = """
def foo():
    try:
        output = 1
    except Exception:
        output = 2
    finally:
        output = 3
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

# ------------------ EDGE TEST CASES ------------------

def test_no_output_assignment():
    # Test function with no assignment to output
    code = """
def foo():
    x = 1
    return x
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_as_attribute():
    # Test assignment to self.output or obj.output should not be counted
    code = """
def foo():
    self.output = 1
    obj.output = 2
    output = 3
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_as_global():
    # Test assignment to global output should be counted if in function body
    code = """
output = 0
def foo():
    global output
    output = 5
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_comprehension():
    # Test assignment in a comprehension (should not be possible)
    code = """
def foo():
    output = 1
    xs = [output for i in range(3)]
    return output
"""
    # Only the direct assignment should be counted
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_deeply_nested_function():
    # Test output assignment in deeply nested function
    code = """
def foo():
    def bar():
        def baz():
            output = 1
    output = 2
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_classmethod_and_staticmethod():
    # Test output assignment in classmethod and staticmethod
    code = """
class C:
    @classmethod
    def foo(cls):
        output = 1
        return output
    @staticmethod
    def bar():
        output = 2
        return output
"""
    codeflash_output = find_codeflash_output_assignments("C.foo", code)
    codeflash_output = find_codeflash_output_assignments("C.bar", code)

def test_function_not_found():
    # Test when the qualified name is not present
    code = """
def foo():
    output = 1
"""
    codeflash_output = find_codeflash_output_assignments("bar", code)

def test_output_in_multiple_functions():
    # Test that only the specified function's assignments are returned
    code = """
def foo():
    output = 1
def bar():
    output = 2
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)
    codeflash_output = find_codeflash_output_assignments("bar", code)

def test_output_in_method_of_nested_class():
    # Test assignment in method of nested class
    code = """
class A:
    class B:
        def foo(self):
            output = 1
            return output
"""
    codeflash_output = find_codeflash_output_assignments("A.B.foo", code)

def test_output_in_function_with_decorator():
    # Test function with a decorator
    code = """
@decorator
def foo():
    output = 1
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_function_with_docstring():
    # Test function with a docstring
    code = '''
def foo():
    """This is a docstring"""
    output = 1
    return output
'''
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_function_with_multiline_assignment():
    # Test assignment split across multiple lines
    code = """
def foo():
    output = (
        1 +
        2
    )
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_function_with_pass():
    # Test function with only pass
    code = """
def foo():
    pass
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_shadowed_in_inner_scope():
    # Assignment to output in an inner function should not be counted
    code = """
def foo():
    output = 1
    def bar():
        nonlocal output
        output = 2
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_output_in_function_with_comment():
    # Test assignment with comment
    code = """
def foo():
    output = 1  # assign output
    return output
"""
    codeflash_output = find_codeflash_output_assignments("foo", code)

# ------------------ LARGE SCALE TEST CASES ------------------

def test_large_function_many_assignments():
    # Test a function with many assignments to output
    code_lines = ["def foo():"] + [f"    output = {i}" for i in range(100)] + ["    return output"]
    code = "\n".join(code_lines)
    expected = list(range(2, 102))
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_large_number_of_functions():
    # Test with many functions, only one of which has output assignments
    code_lines = []
    for i in range(50):
        code_lines.append(f"def func_{i}():")
        code_lines.append(f"    x = {i}")
    code_lines.append("def target():")
    for j in range(100):
        code_lines.append(f"    output = {j}")
    code = "\n".join(code_lines)
    expected = list(range(2 + 2*50, 2 + 2*50 + 100))
    codeflash_output = find_codeflash_output_assignments("target", code)

def test_large_nested_classes_and_methods():
    # Test with nested classes and methods, and output assignment in a deep method
    code_lines = ["class A:"]
    for i in range(10):
        code_lines.append(f"    class B{i}:")
        code_lines.append(f"        def foo(self):")
        code_lines.append(f"            output = {i}")
        code_lines.append(f"            return output")
    code = "\n".join(code_lines)
    for i in range(10):
        qual = f"A.B{i}.foo"
        # The assignment line is 3 + 4*i (class header + lines in each class)
        codeflash_output = find_codeflash_output_assignments(qual, code)

def test_performance_with_many_irrelevant_assignments():
    # Test performance when there are many assignments to other variables
    code_lines = ["def foo():"] + [f"    x{i} = {i}" for i in range(900)] + ["    output = 42", "    return output"]
    code = "\n".join(code_lines)
    codeflash_output = find_codeflash_output_assignments("foo", code)

def test_large_code_with_nested_functions_and_classes():
    # Test a large code with many nested functions and classes, only top-level output assignments count
    code_lines = ["def foo():"]
    for i in range(100):
        code_lines.append(f"    output = {i}")
        code_lines.append(f"    def bar_{i}():")
        code_lines.append(f"        output = -{i}")
        code_lines.append(f"    class Baz_{i}:")
        code_lines.append(f"        def qux(self):")
        code_lines.append(f"            output = {i*2}")
    code = "\n".join(code_lines)
    # Only the first line of each block (output = {i}) should be counted
    expected = [2 + 5*i for i in range(100)]
    codeflash_output = find_codeflash_output_assignments("foo", code)
# 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.edit_generated_tests import \
    find_codeflash_output_assignments

# unit tests

# 1. Basic Test Cases

def test_simple_assignment():
    # Test a simple assignment to a variable
    code = "x = 42"
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_multiple_assignments():
    # Test multiple assignments to the same variable
    code = """
x = 1
x = 2
x = 3
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_augmented_assignment():
    # Test augmented assignment (+=)
    code = """
x = 1
x += 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute():
    # Test assignment to an attribute
    code = """
class A:
    def f(self):
        self.x = 10
        self.x += 1
"""
    codeflash_output = find_codeflash_output_assignments("self.x", code)

def test_assignment_to_nested_attribute():
    # Test assignment to a nested attribute
    code = """
obj.attr.x = 1
obj.attr.x += 2
"""
    codeflash_output = find_codeflash_output_assignments("obj.attr.x", code)

def test_assignment_to_multiple_targets():
    # Test assignment to multiple targets (a = b = c)
    code = "a = b = x = 5"
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_for_loop_target():
    # Test assignment in a for loop target
    code = """
for x in range(3):
    pass
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_with_statement_target():
    # Test assignment in a with statement (with ... as x)
    code = """
with open('foo.txt') as x:
    pass
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_annotated_assignment():
    # Test annotated assignment
    code = """
x: int = 5
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_walrus_operator_assignment():
    # Test assignment with the walrus operator
    code = """
if (x := 3) > 0:
    pass
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

# 2. Edge Test Cases

def test_assignment_to_similar_name():
    # Should not match variables with similar names
    code = """
x = 1
xx = 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute_with_similar_name():
    # Should not match attributes with similar but not exact names
    code = """
obj.x = 1
obj.y = 2
obj.x1 = 3
"""
    codeflash_output = find_codeflash_output_assignments("obj.x", code)

def test_assignment_in_function_scope():
    # Should find assignments inside function scope
    code = """
def foo():
    x = 1
    x += 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_in_class_scope():
    # Should find assignments inside class scope
    code = """
class A:
    x = 1
    def method(self):
        self.x = 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)
    codeflash_output = find_codeflash_output_assignments("self.x", code)

def test_assignment_to_tuple_target():
    # Should find assignment to variable in tuple unpacking
    code = """
x, y = 1, 2
"""
    # Only 'x' is tested
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute_in_tuple():
    # Should find assignment to attribute in tuple unpacking
    code = """
a.x, a.y = 1, 2
"""
    codeflash_output = find_codeflash_output_assignments("a.x", code)

def test_assignment_to_list_target():
    # Should not match assignment to variable in list unpacking
    code = """
[x, y] = [1, 2]
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_subscript():
    # Should not match assignment to subscript
    code = """
x[0] = 1
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute_of_subscript():
    # Should not match assignment to attribute of subscript
    code = """
x[0].y = 1
"""
    codeflash_output = find_codeflash_output_assignments("x.y", code)

def test_assignment_to_global_and_local():
    # Should find both global and local assignments
    code = """
x = 1
def foo():
    x = 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute_of_attribute():
    # Should match only the correct attribute chain
    code = """
a.b.c = 1
a.b = 2
a.b.c.d = 3
"""
    codeflash_output = find_codeflash_output_assignments("a.b.c", code)

def test_assignment_to_attribute_with_function_call():
    # Should not match assignment to attribute of function call
    code = """
f().x = 1
"""
    codeflash_output = find_codeflash_output_assignments("f.x", code)

def test_assignment_to_attribute_with_index():
    # Should not match assignment to attribute of index
    code = """
a[0].x = 1
"""
    codeflash_output = find_codeflash_output_assignments("a.x", code)

def test_assignment_to_attribute_with_different_chain():
    # Should not match different attribute chains
    code = """
a.b.x = 1
a.c.x = 2
"""
    codeflash_output = find_codeflash_output_assignments("a.b.x", code)
    codeflash_output = find_codeflash_output_assignments("a.c.x", code)

def test_assignment_to_attribute_with_similar_names():
    # Should not match attribute chains with similar names
    code = """
a.b.x = 1
a.b.y = 2
a.bx = 3
"""
    codeflash_output = find_codeflash_output_assignments("a.b.x", code)
    codeflash_output = find_codeflash_output_assignments("a.b.y", code)
    codeflash_output = find_codeflash_output_assignments("a.bx", code)

def test_assignment_with_comments_and_blank_lines():
    # Should ignore comments and blank lines, but line numbers should be correct
    code = """

# comment
x = 1  # assign
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_with_unicode_variable():
    # Should match unicode variable names
    code = """
变量 = 1
变量 += 2
"""
    codeflash_output = find_codeflash_output_assignments("变量", code)

def test_assignment_with_nonexistent_variable():
    # Should return empty list if variable is not assigned
    code = """
y = 1
z = 2
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_with_empty_code():
    # Should return empty list for empty code
    code = ""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_with_only_comments():
    # Should return empty list for code with only comments
    code = """
# just a comment
# another comment
"""
    codeflash_output = find_codeflash_output_assignments("x", code)

def test_assignment_to_attribute_with_trailing_whitespace():
    # Should match even with trailing whitespace
    code = "self.x = 1    "
    codeflash_output = find_codeflash_output_assignments("self.x", code)


def test_many_assignments_single_variable():
    # Test performance with many assignments to the same variable
    code = "\n".join(f"x = {i}" for i in range(1000))
    codeflash_output = find_codeflash_output_assignments("x", code); result = codeflash_output

def test_many_assignments_many_variables():
    # Test performance with many variables, only one matches
    code = "\n".join(f"var{i} = {i}" for i in range(1000))
    code += "\nx = 5\n"
    codeflash_output = find_codeflash_output_assignments("x", code); result = codeflash_output

def test_many_attribute_assignments():
    # Test many assignments to the same attribute
    code = "\n".join(f"obj.x = {i}" for i in range(1000))
    codeflash_output = find_codeflash_output_assignments("obj.x", code); result = codeflash_output

def test_large_code_with_mixed_assignments():
    # Test large code with a mix of assignments, only some match
    lines = []
    for i in range(500):
        lines.append(f"x = {i}")
        lines.append(f"y = {i}")
        lines.append(f"obj.x = {i}")
        lines.append(f"obj.y = {i}")
    code = "\n".join(lines)
    codeflash_output = find_codeflash_output_assignments("obj.x", code); result = codeflash_output
    # Should be every 3rd line starting from line 3
    expected = [i for i in range(3, 2001, 4)]

def test_large_code_with_no_matches():
    # Test large code with no matching assignments
    code = "\n".join(f"y{i} = {i}" for i in range(1000))
    codeflash_output = find_codeflash_output_assignments("x", code); result = codeflash_output

def test_large_code_with_nested_attributes():
    # Test large code with nested attribute assignments
    code = "\n".join(f"a.b.c = {i}" for i in range(1000))
    codeflash_output = find_codeflash_output_assignments("a.b.c", code); result = codeflash_output

def test_large_code_with_mixed_scopes():
    # Test large code with assignments in different scopes
    code = "\n".join([
        "x = 0",
        "def f():",
        "    x = 1",
        "    for i in range(10):",
        "        x += i",
        "    return x",
        "class C:",
        "    x = 2",
        "    def m(self):",
        "        self.x = 3",
        "        self.x += 1",
        "        return self.x"
    ] * 80)  # 12 lines * 80 = 960 lines
    codeflash_output = find_codeflash_output_assignments("x", code); result = codeflash_output
    # Each block has lines 1, 3, 5, 8 in its 12-line chunk, offset by 12*block
    expected = []
    for block in range(80):
        offset = block * 12
        expected += [offset+1, offset+3, offset+5, offset+8]

def test_large_code_with_unicode_and_ascii():
    # Test large code with both unicode and ascii variable assignments
    code = "\n".join([
        f"x = {i}" if i % 2 == 0 else f"变量 = {i}" for i in range(1000)
    ])
    codeflash_output = find_codeflash_output_assignments("x", code); result_ascii = codeflash_output
    codeflash_output = find_codeflash_output_assignments("变量", code); result_unicode = codeflash_output
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-find_codeflash_output_assignments-mcmnwrek and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #488 by aseembits93 was closed.