⚡️ Speed up function find_functions_with_return_statement by 26% in PR #1460 (call-graphee)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1460

If you approve this dependent PR, these changes will be merged into the original PR branch call-graphee.

This PR will be automatically closed if the original PR is merged.

---

📄 26% (0.26x) speedup for find_functions_with_return_statement in codeflash/discovery/functions_to_optimize.py

⏱️ Runtime : 12.0 milliseconds → 9.48 milliseconds (best of 46 runs)

📝 Explanation and details

The optimized code achieves a 26% runtime improvement by making the AST traversal in function_has_return_statement more targeted and efficient.

Key Optimization:

The critical change is in how function_has_return_statement traverses the AST when searching for Return nodes:

Original approach:

stack.extend(ast.iter_child_nodes(node))

This visits all child nodes including expressions, names, constants, and other non-statement nodes.

Optimized approach:

for child in ast.iter_child_nodes(node):
    if isinstance(child, ast.stmt):
        stack.append(child)

This only pushes statement nodes onto the stack, since Return is a statement type (ast.stmt).

Why This Is Faster:

1. Reduced Node Traversal: In typical Python functions, there are many more expression nodes (variable references, literals, operators, etc.) than statement nodes. For example, a simple return x + y has 1 Return statement but multiple Name and BinOp expression nodes underneath. The optimization skips all the expression-level nodes.

2. Lower Python Overhead: Fewer nodes in the stack means fewer loop iterations, fewer isinstance checks on non-Return nodes, and less list manipulation overhead.

3. Preserved Correctness: Since Return nodes are always statements in Python's AST (they inherit from ast.stmt), filtering to only statement nodes cannot miss any Return nodes.

Performance Impact by Test Case:

The optimization shows particularly strong gains for:

• Functions without returns (up to 91% faster): Early termination without traversing deep expression trees
• Large codebases (34-41% faster on tests with 1000+ functions): The cumulative effect across many function bodies
• Functions with complex expressions but no returns (82% faster): Avoiding expensive traversal of unused expression subtrees
• Generator functions without explicit returns (64% faster): Skipping yield expression internals

The optimization maintains correctness across all test cases including nested classes, async functions, properties, and various control structures, while delivering consistent runtime improvements.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 61 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import ast  # used to build ASTs from source strings
from pathlib import Path  # used to construct file_path argument

import pytest  # used for our unit tests
from codeflash.discovery.functions_to_optimize import \
    find_functions_with_return_statement
from codeflash.models.function_types import FunctionParent, FunctionToOptimize

def test_empty_module_returns_no_functions():
    # Create an empty AST module from an empty source string
    module = ast.parse("")
    # Call the function under test with a dummy file path
    codeflash_output = find_functions_with_return_statement(module, Path("empty.py")); results = codeflash_output # 3.61μs -> 3.55μs (1.69% faster)

def test_single_function_with_return_is_detected():
    # Build a module with a single top-level function that contains a return
    src = """
def myfunc():
    return 42
"""
    module = ast.parse(src)
    # Use a concrete Path object as required
    file_path = Path("single.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); results = codeflash_output # 18.1μs -> 18.0μs (0.946% faster)
    f: FunctionToOptimize = results[0]

def test_function_without_return_is_ignored():
    # Function that has no return statement should not be included
    src = """
def no_return():
    x = 1
    y = x + 2
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("no_return.py")); results = codeflash_output # 18.5μs -> 9.68μs (91.3% faster)

def test_async_function_with_return_is_detected_and_marked_async():
    # Async function should be detected and flagged via is_async
    src = """
async def async_fn():
    return "ok"
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("async.py")); results = codeflash_output # 17.4μs -> 17.1μs (2.06% faster)
    f = results[0]

def test_method_inside_class_has_class_parent_recorded():
    # Method defined inside a class should produce a FunctionParent entry for the class
    src = """
class MyClass:
    def method(self):
        return 1
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("class_method.py")); results = codeflash_output # 22.0μs -> 21.4μs (2.72% faster)
    f = results[0]
    parent = f.parents[0]

def test_property_decorated_method_is_excluded():
    # A method decorated with @property should be excluded by the filter
    src = """
class C:
    @property
    def p(self):
        return 5
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("property.py")); results = codeflash_output # 16.1μs -> 14.2μs (12.7% faster)

def test_nested_inner_function_return_counts_for_outer_only_and_inner_not_listed():
    # Outer function contains an inner function that has a return.
    # find_functions_with_return_statement should report the outer function (because a Return exists under it)
    # but it should not create a separate entry for the inner function because recursion into function bodies is skipped.
    src = """
def outer():
    def inner():
        return 7
    x = 3
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("nested_inner.py")); results = codeflash_output # 22.6μs -> 21.0μs (7.68% faster)
    # Confirm the inner function is not separately listed
    names = [r.function_name for r in results]

def test_nested_classes_produce_ordered_parents_list():
    # Nested classes: A -> B -> method; parents should reflect both classes in order
    src = """
class A:
    class B:
        def m(self):
            return None
"""
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("nested_classes.py")); results = codeflash_output # 24.9μs -> 24.4μs (2.30% faster)
    f = results[0]

def test_large_scale_many_functions_half_with_return():
    # Construct a module with 1000 functions; even-indexed functions will contain a return.
    # This tests scalability and correctness for many functions.
    N = 1000
    lines = []
    for i in range(N):
        if i % 2 == 0:
            # functions with return
            lines.append(f"def fn_{i}():\n    return {i}\n")
        else:
            # functions without return
            lines.append(f"def fn_{i}():\n    x = {i}\n")
    src = "\n".join(lines)
    module = ast.parse(src)
    codeflash_output = find_functions_with_return_statement(module, Path("many_funcs.py")); results = codeflash_output # 5.94ms -> 4.20ms (41.4% faster)
    # Verify that names correspond to even-indexed functions and there are no duplicates
    found_names = {r.function_name for r in results}
    expected_names = {f"fn_{i}" for i in range(0, N, 2)}

import ast
from pathlib import Path

# imports
import pytest
from codeflash.discovery.functions_to_optimize import \
    find_functions_with_return_statement
from codeflash.models.function_types import FunctionParent, FunctionToOptimize

def test_single_function_with_return_statement():
    """Test finding a simple function with a return statement."""
    code = """
def simple_func():
    return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 22.4μs -> 25.8μs (13.3% slower)

def test_single_function_without_return_statement():
    """Test that a function without return statement is not included."""
    code = """
def no_return_func():
    x = 1
    print(x)
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.9μs -> 9.80μs (82.5% faster)

def test_async_function_with_return():
    """Test finding an async function with a return statement."""
    code = """
async def async_func():
    return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 18.2μs -> 17.8μs (2.66% faster)

def test_function_with_implicit_return():
    """Test that function without explicit return is not included."""
    code = """
def implicit_return():
    pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 10.0μs -> 7.34μs (36.4% faster)

def test_class_with_method_with_return():
    """Test finding a method in a class that has a return statement."""
    code = """
class MyClass:
    def method_with_return(self):
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 23.1μs -> 22.8μs (1.19% faster)

def test_multiple_functions_mixed_returns():
    """Test multiple functions where some have returns and some don't."""
    code = """
def func_with_return():
    return 42

def func_without_return():
    x = 1

def another_with_return():
    return "hello"
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 33.7μs -> 28.2μs (19.5% faster)
    names = {r.function_name for r in result}

def test_nested_classes_with_methods():
    """Test finding methods in nested classes."""
    code = """
class Outer:
    class Inner:
        def method_with_return(self):
            return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 25.1μs -> 24.6μs (1.71% faster)

def test_property_decorator_excluded():
    """Test that functions decorated with @property are excluded."""
    code = """
class MyClass:
    @property
    def my_property(self):
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 16.1μs -> 14.5μs (11.3% faster)

def test_property_decorator_with_other_decorators():
    """Test that @property decorator is detected even with other decorators."""
    code = """
class MyClass:
    @some_decorator
    @property
    def my_property(self):
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.1μs -> 14.5μs (18.2% faster)

def test_function_with_multiple_returns():
    """Test function with multiple return statements."""
    code = """
def multi_return(x):
    if x > 0:
        return 1
    else:
        return 0
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 19.3μs -> 19.1μs (0.891% faster)

def test_function_with_return_in_loop():
    """Test function with return statement inside a loop."""
    code = """
def loop_return(items):
    for item in items:
        if item > 10:
            return item
    return None
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.6μs -> 17.2μs (2.44% faster)

def test_function_with_return_in_try_except():
    """Test function with return inside try-except block."""
    code = """
def try_return():
    try:
        return 42
    except Exception:
        pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 22.5μs -> 18.5μs (21.6% faster)

def test_function_with_nested_function_with_return():
    """Test that nested functions are not recursed into."""
    code = """
def outer():
    def inner():
        return 42
    pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 19.7μs -> 20.1μs (2.04% slower)

def test_function_with_lambda_containing_return():
    """Test function containing lambda (lambdas don't use return, but func might)."""
    code = """
def with_lambda():
    f = lambda x: x * 2
    return f
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.3μs -> 17.2μs (0.523% faster)

def test_method_overriding_in_subclass():
    """Test multiple classes with methods having same name."""
    code = """
class ClassA:
    def method(self):
        return 1

class ClassB:
    def method(self):
        pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 34.4μs -> 29.4μs (17.0% faster)

def test_empty_module():
    """Test with an empty AST module."""
    code = ""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 3.27μs -> 3.18μs (2.83% faster)

def test_module_with_only_classes_no_methods():
    """Test module with class but no methods."""
    code = """
class EmptyClass:
    pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 12.1μs -> 12.0μs (0.500% faster)

def test_function_with_bare_return():
    """Test function with bare return statement (returns None)."""
    code = """
def bare_return():
    if True:
        return
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 19.1μs -> 18.8μs (1.65% faster)

def test_very_deeply_nested_classes():
    """Test functions in deeply nested class hierarchy."""
    code = """
class L1:
    class L2:
        class L3:
            class L4:
                def deep_method(self):
                    return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 30.0μs -> 29.6μs (1.29% faster)

def test_function_with_return_of_complex_expression():
    """Test function returning complex expression."""
    code = """
def complex_return():
    return (1 + 2) * [1, 2, 3] or {"key": "value"}
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.3μs -> 17.0μs (1.76% faster)

def test_function_with_return_in_nested_if():
    """Test return in deeply nested if statements."""
    code = """
def nested_if_return():
    if True:
        if True:
            if True:
                return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 20.4μs -> 20.0μs (2.10% faster)

def test_function_with_return_in_with_statement():
    """Test return inside with statement."""
    code = """
def with_return():
    with open("file") as f:
        return f.read()
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 18.5μs -> 18.5μs (0.167% faster)

def test_function_with_return_in_finally():
    """Test return in finally block."""
    code = """
def finally_return():
    try:
        x = 1
    finally:
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 18.8μs -> 18.7μs (0.910% faster)

def test_static_method_with_return():
    """Test static method with return statement."""
    code = """
class MyClass:
    @staticmethod
    def static_method():
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 23.9μs -> 22.1μs (8.02% faster)

def test_class_method_with_return():
    """Test class method with return statement."""
    code = """
class MyClass:
    @classmethod
    def class_method(cls):
        return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 23.3μs -> 21.6μs (7.83% faster)

def test_dunder_methods_with_return():
    """Test special/dunder methods with return."""
    code = """
class MyClass:
    def __init__(self):
        self.x = 1
    
    def __str__(self):
        return str(self.x)
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 33.5μs -> 25.4μs (32.1% faster)

def test_function_with_yield_but_no_return():
    """Test generator function with yield but no explicit return."""
    code = """
def generator():
    yield 1
    yield 2
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 13.5μs -> 8.23μs (63.5% faster)

def test_function_with_yield_and_return():
    """Test generator function with both yield and return."""
    code = """
def generator_with_return():
    yield 1
    return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.8μs -> 17.8μs (0.000% faster)

def test_async_generator_with_return():
    """Test async generator with return."""
    code = """
async def async_gen():
    yield 1
    return 42
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.5μs -> 17.0μs (2.53% faster)

def test_special_file_path():
    """Test with various file paths."""
    code = """
def func():
    return 1
"""
    module = ast.parse(code)
    
    # Test with different path formats
    paths = [
        Path("/absolute/path/to/file.py"),
        Path("relative/path/file.py"),
        Path("/path/with spaces/file.py"),
        Path("/path/with-dashes/file.py"),
    ]
    
    for file_path in paths:
        codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 43.3μs -> 41.1μs (5.36% faster)

def test_function_name_preservation():
    """Test that function names are preserved exactly."""
    code = """
def _private_func():
    return 1

def __dunder_func__():
    return 2

def CamelCaseFunc():
    return 3
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 30.7μs -> 30.5μs (0.655% faster)
    names = {r.function_name for r in result}

def test_return_with_no_value():
    """Test bare return statements."""
    code = """
def bare_return():
    if condition:
        return
    print("fallthrough")
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 23.0μs -> 19.6μs (17.4% faster)

def test_class_with_init_without_return():
    """Test __init__ without explicit return (implicit None)."""
    code = """
class MyClass:
    def __init__(self):
        self.x = 1
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 23.0μs -> 15.2μs (51.1% faster)

def test_parent_path_preservation():
    """Test that parent path is correctly preserved for methods."""
    code = """
class OuterClass:
    class InnerClass:
        def target_method(self):
            return 1
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 25.2μs -> 24.4μs (3.24% faster)

def test_return_statement_in_comprehension():
    """Test that return inside comprehension context is found."""
    code = """
def with_return():
    result = [x for x in range(10)]
    return result
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.4μs -> 17.5μs (0.635% slower)

def test_multiple_classes_at_module_level():
    """Test multiple classes each with methods."""
    code = """
class Class1:
    def method1(self):
        return 1

class Class2:
    def method2(self):
        return 2

class Class3:
    def method3(self):
        pass
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 43.8μs -> 39.0μs (12.5% faster)

def test_function_with_walrus_operator_return():
    """Test function with walrus operator in return."""
    code = """
def walrus_return():
    return (x := 42)
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 17.1μs -> 16.8μs (1.97% faster)

def test_method_with_type_hints_and_return():
    """Test function with type hints."""
    code = """
def typed_func(x: int) -> int:
    return x + 1
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 18.7μs -> 16.7μs (11.9% faster)

def test_async_function_without_return():
    """Test async function without return."""
    code = """
async def async_no_return():
    await something()
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 13.7μs -> 7.18μs (91.2% faster)

def test_many_functions_at_module_level():
    """Test module with many functions, some with returns."""
    # Create 100 functions, 50 with returns and 50 without
    code_lines = []
    for i in range(50):
        code_lines.append(f"""
def func_with_return_{i}():
    return {i}
""")
    for i in range(50):
        code_lines.append(f"""
def func_without_return_{i}():
    x = {i}
""")
    
    code = "".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 599μs -> 429μs (39.6% faster)

def test_many_methods_in_single_class():
    """Test single class with many methods."""
    code_lines = ["class LargeClass:"]
    for i in range(100):
        code_lines.append(f"""
    def method_{i}(self):
        return {i}
""")
    
    code = "".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 566μs -> 581μs (2.71% slower)

def test_large_mixed_class_structure():
    """Test complex structure with many classes and methods."""
    code_lines = []
    for c in range(20):
        code_lines.append(f"class Class{c}:")
        for m in range(10):
            if m % 2 == 0:
                code_lines.append(f"""
    def method_{m}(self):
        return {m}
""")
            else:
                code_lines.append(f"""
    def method_{m}(self):
        pass
""")
    
    code = "".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 1.14ms -> 853μs (34.0% faster)

def test_function_with_many_nested_control_structures():
    """Test function with deeply nested control flow containing return."""
    code_lines = ["def complex_func():"]
    code_lines.append("    if True:")
    for i in range(50):
        code_lines.append(" " * (8 + 4 * i) + "if True:")
    code_lines.append(" " * (8 + 4 * 50) + "return 42")
    
    code = "\n".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 66.7μs -> 68.6μs (2.82% slower)

def test_many_decorators_on_functions():
    """Test functions with many decorators (some with property)."""
    code_lines = []
    for i in range(50):
        code_lines.append(f"""
@decorator1
@decorator2
def func_{i}():
    return {i}
""")
    for i in range(50, 100):
        code_lines.append(f"""
@decorator1
@property
@decorator2
def func_{i}():
    return {i}
""")
    
    code = "".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 679μs -> 437μs (55.4% faster)

def test_return_in_various_container_types():
    """Test return statements in functions with various container literals."""
    code_lines = []
    for i in range(100):
        if i % 5 == 0:
            code_lines.append(f"def func_{i}():\n    return [1, 2, 3]")
        elif i % 5 == 1:
            code_lines.append(f"def func_{i}():\n    return {{1, 2, 3}}")
        elif i % 5 == 2:
            code_lines.append(f"def func_{i}():\n    return {{'a': 1}}")
        elif i % 5 == 3:
            code_lines.append(f"def func_{i}():\n    return (1, 2, 3)")
        else:
            code_lines.append(f"def func_{i}():\n    return 42")
    
    code = "\n".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 553μs -> 564μs (1.91% slower)

def test_mixed_async_and_sync_functions_large():
    """Test module with many async and sync functions."""
    code_lines = []
    for i in range(100):
        if i % 2 == 0:
            code_lines.append(f"async def async_func_{i}():\n    return {i}")
        else:
            code_lines.append(f"def sync_func_{i}():\n    return {i}")
    
    code = "\n".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 552μs -> 563μs (2.01% slower)
    async_count = sum(1 for r in result if r.is_async)
    sync_count = sum(1 for r in result if not r.is_async)

def test_large_class_inheritance_like_structure():
    """Test many classes each with their own methods."""
    code_lines = []
    for c in range(100):
        code_lines.append(f"class Class{c}:")
        code_lines.append(f"    def method(self):\n        return {c}")
    
    code = "\n".join(code_lines)
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 817μs -> 829μs (1.47% slower)

def test_result_list_independence():
    """Test that result objects are independent instances."""
    code = """
def func1():
    return 1

def func2():
    return 2
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 25.4μs -> 25.1μs (1.28% faster)

def test_parents_list_is_copy():
    """Test that parents list is independent across results."""
    code = """
class MyClass:
    def method1(self):
        return 1
    def method2(self):
        return 2
"""
    module = ast.parse(code)
    file_path = Path("/test/module.py")
    codeflash_output = find_functions_with_return_statement(module, file_path); result = codeflash_output # 29.6μs -> 28.8μs (2.65% faster)
# 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-pr1460-2026-02-18T22.22.36 and push.

[claude]

PR Review Summary

Prek Checks

✅ All checks pass — ruff check and ruff format both clean. No auto-fixes needed.

Mypy

✅ No new mypy errors introduced by this PR. All mypy errors in the output are pre-existing in files pulled in transitively (e.g., console.py, javascript/support.py).

Code Review

No critical bugs, security vulnerabilities, or breaking API changes found.

This PR introduces a Jedi-based CallGraph for Python dependency resolution, replaces the Jedi Name object on FunctionSource with a simple definition_type: str | None, refactors FunctionWithReturnStatement to an iterative DFS function, and adds a DependencyResolver protocol. Key observations:

• The FunctionSource model no longer requires arbitrary_types_allowed — good decoupling from Jedi internals
• find_functions_with_return_statement iterative DFS correctly preserves original visitor behavior (skips recursion into function bodies, tracks only ClassDef parents)
• function_has_return_statement optimization to filter ast.stmt children is correct since ast.Return is a subclass of ast.stmt
• SQL in CallGraph uses parameterized queries throughout — no injection risk
• compat.py simplification from lazy @property to module-level variables is fine (the directory creation at import time is harmless)
• Call graph is correctly skipped in CI (is_ci() check) and the resolver is properly closed in a finally block
• Token limit increase from 48000 → 64000 is intentional

Test Coverage

File	PR	Main	Δ
cli_cmds/console.py	26%	48%	⚠️ -22%
code_utils/code_replacer.py	83%	83%	—
code_utils/compat.py	100%	97%	+3%
code_utils/config_consts.py	88%	88%	—
discovery/functions_to_optimize.py	68%	69%	-1%
languages/__init__.py	30%	30%	—
languages/base.py	98%	99%	-1%
languages/javascript/support.py	74%	74%	—
languages/python/__init__.py	100%	100%	—
languages/python/call_graph.py (new)	72%	N/A	✅
languages/python/context/code_context_extractor.py	92%	91%	+1%
languages/python/context/unused_definition_remover.py	94%	94%	—
languages/python/support.py	51%	51%	—
models/models.py	78%	78%	—
optimization/function_optimizer.py	19%	18%	+1%
optimization/optimizer.py	19%	20%	-1%
Total (changed files)	61%	62%	-1%

Notes:

• call_graph.py (new, 72%): Above the 75% threshold by a small margin — the untested code is primarily the parallel indexing path (ProcessPoolExecutor) and error handling branches. The 475-line test_call_graph.py provides good unit test coverage.
• console.py (-22%): The new Rich live display functions (call_graph_live_display, call_graph_summary) are not exercised by unit tests. This is UI/display code that's difficult to test in isolation.
• Overall coverage: The -1% total change is negligible and driven by new UI code in console.py.
• 8 test failures in test_tracer.py are pre-existing on both branches and unrelated to this PR.

---

Last updated: 2026-02-18T22:50Z