⚡️ Speed up function _extract_calling_function by 18% in PR #1256 (refactor/use-function-to-optimize-in-js)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1256

If you approve this dependent PR, these changes will be merged into the original PR branch refactor/use-function-to-optimize-in-js.

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

---

📄 18% (0.18x) speedup for _extract_calling_function in codeflash/code_utils/code_extractor.py

⏱️ Runtime : 6.31 milliseconds → 5.34 milliseconds (best of 224 runs)

📝 Explanation and details

The optimized code achieves an 18% runtime improvement by replacing the exhaustive AST traversal with a pruned depth-first search (DFS) that skips irrelevant subtrees.

Key optimization: Spatial pruning in AST traversal

The original code uses ast.walk(tree), which visits every node in the AST regardless of whether it could possibly contain the target function. The optimized version implements a stack-based DFS that prunes entire subtrees when their line ranges don't contain ref_line:

# Prune nodes that can't contain the reference line
if node_start is not None:
    node_end = getattr(node, "end_lineno", node_start) or node_start
    if not (node_start <= ref_line <= node_end):
        continue  # Skip this entire subtree

This spatial pruning is highly effective because:

• ASTs for typical Python files contain hundreds of nodes (as seen in the profiler: 781 nodes visited in the original)
• Most nodes fall outside the reference line's range
• The optimized version visits only ~336 nodes (57% reduction), as shown in the profiler's while stack iteration count

Secondary optimization: Deferred string splitting

The code also moves source_code.splitlines() to execute only after finding a matching function, avoiding unnecessary work when no match exists (13 of 50 test cases in the profiler).

Performance characteristics based on test results:

The optimization is most effective for:

• Large files with many functions (83.4% speedup for 100-function file, 52.7% speedup for mixed-content file)
• Edge cases with extreme ref_line values (43-46% speedup when ref_line is far outside valid ranges, enabling early pruning)
• Ref_line outside function ranges (20-28% speedup when ref_line doesn't match any function)

The speedup is more modest (10-20%) for simple cases where most nodes need visiting anyway, but these cases still benefit from reduced overhead.

Line profiler data confirms the optimization: the original code spent 55.8% of time in ast.walk() iteration, while the optimized version eliminates this bottleneck through intelligent pruning. The pruning logic itself adds only 5.8% overhead (lines checking node_start and bounds), which is vastly outweighed by the reduction in nodes processed.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from types import SimpleNamespace

import codeflash.languages.treesitter_utils as ts_utils
# imports
import pytest  # used for our unit tests
from codeflash.code_utils.code_extractor import _extract_calling_function
from codeflash.languages.base import Language
from codeflash.languages.treesitter_utils import TreeSitterLanguage

# Helper to pick a non-PYTHON language value from Language enum if available,
# otherwise fall back to a plain object (which will still trigger the JS branch
# because it will not be equal to Language.PYTHON).
def _non_python_language():
    try:
        # Try to find any member of Language that is not PYTHON
        for member in Language:
            if member != Language.PYTHON:
                return member
    except Exception:
        # If Language is not iterable or something unexpected happens, fall back.
        return object()
    # If only PYTHON exists, still return a plain object to trigger JS branch
    return object()

def test_python_basic_extraction():
    # Simple source with two top-level functions. We will extract 'target' given a
    # ref_line located inside it.
    source = (
        "def a():\n"
        "    x = 1\n"
        "\n"
        "def target():\n"
        "    y = 2\n"
        "    z = 3\n"
        "\n"
        "other = 0\n"
    )
    # Line numbers (1-based):
    # 1 def a():
    # 2     x = 1
    # 3
    # 4 def target():
    # 5     y = 2
    # 6     z = 3
    # 7
    # 8 other = 0
    ref_line = 5  # inside 'target'
    codeflash_output = _extract_calling_function(source, "target", ref_line, Language.PYTHON); result = codeflash_output # 47.4μs -> 40.1μs (18.1% faster)
    # Expect the exact source lines for 'target' only (lines 4-6).
    expected = "def target():\n    y = 2\n    z = 3"

def test_python_ref_outside_returns_none():
    # Same source as before; pick a ref_line outside any function.
    source = (
        "def foo():\n"
        "    pass\n"
        "\n"
        "bar = 1\n"
    )
    # 'bar = 1' is line 4 which is outside any function -> expect None
    codeflash_output = _extract_calling_function(source, "foo", 4, Language.PYTHON); result = codeflash_output # 37.0μs -> 30.7μs (20.4% faster)

def test_python_nested_and_async_function_extraction():
    # Test nested functions and async function extraction.
    source_lines = []
    source_lines.append("def outer():")
    source_lines.append("    def inner():")
    source_lines.append("        a = 1")
    source_lines.append("    return inner")
    source_lines.append("")  # blank line
    source_lines.append("async def async_fn():")
    source_lines.append("    await something()")
    source = "\n".join(source_lines) + "\n"
    # Locate 'inner' which starts at line 2 and body at 3
    ref_line_inner = 3
    codeflash_output = _extract_calling_function(source, "inner", ref_line_inner, Language.PYTHON); res_inner = codeflash_output # 54.0μs -> 45.6μs (18.4% faster)
    # Expect only the inner function definition (lines 2-3)
    expected_inner = "    def inner():\n        a = 1"

    # Now test async function extraction: async_fn starts at line 6 (1-based)
    ref_line_async = 7  # inside await line
    codeflash_output = _extract_calling_function(source, "async_fn", ref_line_async, Language.PYTHON); res_async = codeflash_output # 33.0μs -> 27.3μs (20.9% faster)
    # Expected to return the async function lines (lines 6-7)
    expected_async = "async def async_fn():\n    await something()"

def test_malformed_python_returns_none():
    # Provide syntactically invalid Python code. The implementation catches
    # exceptions and should return None rather than raising.
    bad_source = "def broken(:\n"
    codeflash_output = _extract_calling_function(bad_source, "broken", 1, Language.PYTHON); res = codeflash_output # 26.7μs -> 26.6μs (0.413% faster)

def test_python_function_name_present_but_ref_outside_returns_none():
    # Function exists but the reference line is outside the function range.
    source = (
        "def myfunc():\n"
        "    x = 10\n"
        "\n"
        "# reference is below\n"
        "y = 5\n"
    )
    # myfunc covers lines 1-2; ref_line 5 is outside -> None
    codeflash_output = _extract_calling_function(source, "myfunc", 5, Language.PYTHON); res = codeflash_output # 42.2μs -> 33.0μs (27.9% faster)

def test_large_scale_python_many_functions_performance():
    # Build a large file with many small functions (but keep under 1000 items).
    # We create 200 functions and target one in the middle to ensure scalability.
    num_funcs = 200
    lines = []
    # Generate functions: each has two lines (def and a body), plus a blank line.
    for i in range(num_funcs):
        lines.append(f"def f_{i}():")
        lines.append(f"    x = {i}")
        lines.append("")  # blank line
    source = "\n".join(lines) + "\n"
    # Choose target index and compute approximate ref_line.
    target_idx = 150
    # Each function takes 3 lines, so start_line = 1 + target_idx * 3
    start_line = 1 + target_idx * 3
    # ref_line pick the body line (start_line + 1)
    ref_line = start_line + 1
    function_name = f"f_{target_idx}"
    codeflash_output = _extract_calling_function(source, function_name, ref_line, Language.PYTHON); res = codeflash_output # 1.35ms -> 1.16ms (16.0% faster)
    # Expected two non-blank lines for the function (def and body)
    expected = f"def {function_name}():\n    x = {target_idx}"

def test_js_basic_extraction(monkeypatch):
    # Monkeypatch the TreeSitterAnalyzer to return a function descriptor that
    # matches the query. Use SimpleNamespace objects to avoid defining new classes.
    non_py_lang = _non_python_language()

    # Fake analyzer that always returns a matching function
    class FakeAnalyzer:
        def __init__(self, language):
            # store language to allow assertions if needed
            self.language = language

        def find_functions(self, source, include_methods=True):
            # Return a single "function" that spans lines 2-4 and matches name 'foo'
            fn = SimpleNamespace(
                name="foo",
                start_line=2,
                end_line=4,
                source_text="function foo() {\n  return 42;\n}"
            )
            return [fn]

    # Patch the real TreeSitterAnalyzer with our FakeAnalyzer
    monkeypatch.setattr(ts_utils, "TreeSitterAnalyzer", FakeAnalyzer)

    # Source text is irrelevant for the fake analyzer but kept for clarity.
    source = "/* dummy js file */\nfunction foo() {\n  return 42;\n}\n"
    # ref_line = 3 is within the function (2-4)
    codeflash_output = _extract_calling_function(source, "foo", 3, non_py_lang); res = codeflash_output # 7.08μs -> 6.40μs (10.6% faster)

def test_js_fallback_language_try_until_success(monkeypatch):
    # Ensure that if the analyzer raises for the first languages, the code will
    # continue and succeed for a later language.
    non_py_lang = _non_python_language()
    TreeSitterLanguage = ts_utils.TreeSitterLanguage

    # Build a fake analyzer factory that raises for TYPESCRIPT and TSX but works for JAVASCRIPT.
    class FailingThenWorkingAnalyzer:
        def __init__(self, language):
            # language is an enum value; raise for two first languages
            if language in (TreeSitterLanguage.TYPESCRIPT, TreeSitterLanguage.TSX):
                raise RuntimeError("Simulated init failure for this language")
            # otherwise ok
            self.language = language

        def find_functions(self, source, include_methods=True):
            # Return a function only if language is JAVASCRIPT
            if self.language == TreeSitterLanguage.JAVASCRIPT:
                return [
                    SimpleNamespace(
                        name="bar",
                        start_line=1,
                        end_line=3,
                        source_text="function bar() {\n  // body\n}"
                    )
                ]
            # Otherwise, simulate no functions found
            return []

    monkeypatch.setattr(ts_utils, "TreeSitterAnalyzer", FailingThenWorkingAnalyzer)

    source = "function bar() {\n  // body\n}\n"
    # ref_line 2 is inside the function
    codeflash_output = _extract_calling_function(source, "bar", 2, non_py_lang); res = codeflash_output # 10.1μs -> 9.27μs (8.98% faster)

def test_js_no_matching_function_returns_none(monkeypatch):
    # Analyzer returns functions but none match the requested name -> expect None
    non_py_lang = _non_python_language()

    class AnalyzerNoMatch:
        def __init__(self, language):
            self.language = language

        def find_functions(self, source, include_methods=True):
            return [
                SimpleNamespace(name="something_else", start_line=1, end_line=10, source_text="function something_else(){}")
            ]

    monkeypatch.setattr(ts_utils, "TreeSitterAnalyzer", AnalyzerNoMatch)

    source = "function something_else(){}\n"
    codeflash_output = _extract_calling_function(source, "nope", 2, non_py_lang); res = codeflash_output # 5.70μs -> 5.42μs (5.18% faster)

def test_js_find_functions_raises_but_next_language_handles(monkeypatch):
    # Simulate the case where find_functions raises an exception for one language,
    # but a subsequent language returns valid functions.
    TreeSitterLanguage = ts_utils.TreeSitterLanguage
    non_py_lang = _non_python_language()

    class AnalyzerSometimesBroken:
        def __init__(self, language):
            self.language = language

        def find_functions(self, source, include_methods=True):
            # Raise for TYPESCRIPT, but work for others
            if self.language == TreeSitterLanguage.TYPESCRIPT:
                raise RuntimeError("broken parse")
            return [
                SimpleNamespace(name="x", start_line=1, end_line=2, source_text="function x(){}")
            ]

    monkeypatch.setattr(ts_utils, "TreeSitterAnalyzer", AnalyzerSometimesBroken)

    source = "function x(){}\n"
    # ref_line 1 is inside the function
    codeflash_output = _extract_calling_function(source, "x", 1, non_py_lang); res = codeflash_output # 8.40μs -> 8.05μs (4.35% faster)

def test_js_entire_js_branch_exception_handled(monkeypatch):
    # If importing or something unexpected raises an exception inside the JS
    # branch, the function should catch it and return None rather than raising.
    # To simulate an import-time exception inside the function's try/except, we
    # monkeypatch the module attribute the function imports to raise when accessed.
    non_py_lang = _non_python_language()

    # Simulate that accessing TreeSitterAnalyzer raises an unexpected Exception.
    def raise_on_access(name):
        raise RuntimeError("simulated import-time failure")

    # Monkeypatch the module attribute to a callable that raises when invoked.
    monkeypatch.setattr(ts_utils, "TreeSitterAnalyzer", raise_on_access)

    # The function should handle this and return None
    codeflash_output = _extract_calling_function("irrelevant", "anything", 1, non_py_lang); res = codeflash_output # 5.73μs -> 5.58μs (2.69% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest
from codeflash.code_utils.code_extractor import _extract_calling_function
from codeflash.languages.base import Language

def test_python_simple_function_extraction():
    """Test extracting a simple Python function when ref_line is within function."""
    source_code = """def hello():
    print("hello")
    return 42
"""
    codeflash_output = _extract_calling_function(source_code, "hello", 2, Language.PYTHON); result = codeflash_output # 42.9μs -> 36.8μs (16.7% faster)

def test_python_function_with_multiple_lines():
    """Test extracting a Python function that spans multiple lines."""
    source_code = """def calculate(x, y):
    temp = x + y
    result = temp * 2
    return result
"""
    codeflash_output = _extract_calling_function(source_code, "calculate", 3, Language.PYTHON); result = codeflash_output # 47.9μs -> 42.2μs (13.5% faster)

def test_python_async_function_extraction():
    """Test extracting an async Python function."""
    source_code = """async def fetch_data():
    await some_operation()
    return data
"""
    codeflash_output = _extract_calling_function(source_code, "fetch_data", 2, Language.PYTHON); result = codeflash_output # 37.4μs -> 32.0μs (17.0% faster)

def test_python_function_with_parameters():
    """Test extracting a Python function with multiple parameters."""
    source_code = """def add(a, b, c=10, *args, **kwargs):
    return a + b + c
"""
    codeflash_output = _extract_calling_function(source_code, "add", 2, Language.PYTHON); result = codeflash_output # 42.6μs -> 36.8μs (15.9% faster)

def test_python_function_not_found():
    """Test that None is returned when function name doesn't exist."""
    source_code = """def hello():
    return 42
"""
    codeflash_output = _extract_calling_function(source_code, "nonexistent", 1, Language.PYTHON); result = codeflash_output # 31.8μs -> 27.9μs (14.0% faster)

def test_python_ref_line_outside_function():
    """Test that None is returned when ref_line is outside all functions."""
    source_code = """def hello():
    return 42

# comment on line 4
x = 10
"""
    codeflash_output = _extract_calling_function(source_code, "hello", 5, Language.PYTHON); result = codeflash_output # 40.3μs -> 32.9μs (22.4% faster)

def test_python_multiple_functions_correct_one_extracted():
    """Test that the correct function is extracted when multiple functions exist."""
    source_code = """def first():
    return 1

def second():
    return 2

def third():
    return 3
"""
    codeflash_output = _extract_calling_function(source_code, "second", 5, Language.PYTHON); result = codeflash_output # 43.5μs -> 36.8μs (18.3% faster)

def test_python_nested_functions():
    """Test extracting an outer function when ref_line is in nested function."""
    source_code = """def outer():
    def inner():
        return 1
    return 2
"""
    # ref_line = 3 is inside inner() but we ask for outer()
    # The function should return outer() since ref_line is within outer's range
    codeflash_output = _extract_calling_function(source_code, "outer", 3, Language.PYTHON); result = codeflash_output # 35.0μs -> 29.3μs (19.5% faster)

def test_python_function_with_docstring():
    """Test extracting a function with a docstring."""
    source_code = '''def documented():
    """This is a docstring."""
    return 42
'''
    codeflash_output = _extract_calling_function(source_code, "documented", 2, Language.PYTHON); result = codeflash_output # 31.5μs -> 26.6μs (18.6% faster)

def test_python_function_with_decorators():
    """Test extracting a decorated Python function."""
    source_code = """@decorator
def decorated():
    return 42
"""
    # Note: ast.parse includes decorators in the function node
    # ref_line = 2 is the def line
    codeflash_output = _extract_calling_function(source_code, "decorated", 2, Language.PYTHON); result = codeflash_output # 31.6μs -> 26.0μs (21.7% faster)

def test_python_empty_source_code():
    """Test with empty source code."""
    codeflash_output = _extract_calling_function("", "any_func", 1, Language.PYTHON); result = codeflash_output # 12.9μs -> 10.3μs (25.4% faster)

def test_python_invalid_syntax():
    """Test with invalid Python syntax - should return None without raising."""
    source_code = """def hello(
    incomplete function
"""
    codeflash_output = _extract_calling_function(source_code, "hello", 1, Language.PYTHON); result = codeflash_output # 42.3μs -> 42.3μs (0.071% slower)

def test_python_ref_line_at_function_start():
    """Test when ref_line is at the start of the function definition."""
    source_code = """def test_func():
    return 42
"""
    codeflash_output = _extract_calling_function(source_code, "test_func", 1, Language.PYTHON); result = codeflash_output # 30.9μs -> 25.0μs (23.2% faster)

def test_python_ref_line_at_function_end():
    """Test when ref_line is at the last line of the function."""
    source_code = """def test_func():
    x = 1
    return x
"""
    codeflash_output = _extract_calling_function(source_code, "test_func", 3, Language.PYTHON); result = codeflash_output # 34.6μs -> 28.9μs (19.6% faster)

def test_python_single_line_function():
    """Test extracting a function defined on a single line (edge case)."""
    source_code = "def oneliner(): return 42"
    codeflash_output = _extract_calling_function(source_code, "oneliner", 1, Language.PYTHON); result = codeflash_output # 27.9μs -> 22.2μs (26.1% faster)

def test_python_function_with_no_body():
    """Test function with only pass statement."""
    source_code = """def empty():
    pass
"""
    codeflash_output = _extract_calling_function(source_code, "empty", 1, Language.PYTHON); result = codeflash_output # 26.1μs -> 20.9μs (25.1% faster)

def test_python_ref_line_zero():
    """Test with ref_line of 0 (edge case)."""
    source_code = """def test():
    return 1
"""
    codeflash_output = _extract_calling_function(source_code, "test", 0, Language.PYTHON); result = codeflash_output # 31.4μs -> 21.6μs (45.1% faster)

def test_python_ref_line_negative():
    """Test with negative ref_line (edge case)."""
    source_code = """def test():
    return 1
"""
    codeflash_output = _extract_calling_function(source_code, "test", -1, Language.PYTHON); result = codeflash_output # 30.6μs -> 21.4μs (43.3% faster)

def test_python_ref_line_very_large():
    """Test with ref_line larger than file."""
    source_code = """def test():
    return 1
"""
    codeflash_output = _extract_calling_function(source_code, "test", 9999, Language.PYTHON); result = codeflash_output # 30.5μs -> 20.9μs (45.9% faster)

def test_python_function_name_is_empty_string():
    """Test with empty string as function_name."""
    source_code = """def test():
    return 1
"""
    codeflash_output = _extract_calling_function(source_code, "", 1, Language.PYTHON); result = codeflash_output # 30.1μs -> 27.0μs (11.4% faster)

def test_python_function_name_case_sensitive():
    """Test that function name matching is case-sensitive."""
    source_code = """def MyFunc():
    return 1
"""
    codeflash_output = _extract_calling_function(source_code, "myfunc", 1, Language.PYTHON); result = codeflash_output # 29.7μs -> 26.5μs (12.1% faster)
    
    codeflash_output = _extract_calling_function(source_code, "MyFunc", 1, Language.PYTHON); result = codeflash_output # 18.3μs -> 14.5μs (26.2% faster)

def test_python_function_with_complex_decorators():
    """Test function with multiple complex decorators."""
    source_code = """@decorator1
@decorator2(arg="value")
@decorator3
def complex_decorated():
    return 42
"""
    codeflash_output = _extract_calling_function(source_code, "complex_decorated", 4, Language.PYTHON); result = codeflash_output # 44.1μs -> 38.8μs (13.8% faster)

def test_python_function_with_class_definition_inside():
    """Test function containing a class definition."""
    source_code = """def factory():
    class Inner:
        pass
    return Inner()
"""
    codeflash_output = _extract_calling_function(source_code, "factory", 2, Language.PYTHON); result = codeflash_output # 36.9μs -> 31.4μs (17.6% faster)

def test_python_function_with_lambda():
    """Test function containing lambda expressions."""
    source_code = """def with_lambda():
    f = lambda x: x * 2
    return f(5)
"""
    codeflash_output = _extract_calling_function(source_code, "with_lambda", 2, Language.PYTHON); result = codeflash_output # 44.2μs -> 38.9μs (13.6% faster)

def test_python_whitespace_only_source():
    """Test with source code containing only whitespace."""
    source_code = "   \n  \n   \n"
    codeflash_output = _extract_calling_function(source_code, "any_func", 1, Language.PYTHON); result = codeflash_output # 13.2μs -> 10.2μs (29.8% faster)

def test_python_function_with_multiline_string():
    """Test function containing multiline strings."""
    source_code = '''def with_string():
    text = """
    This is a multiline
    string that spans
    multiple lines
    """
    return text
'''
    codeflash_output = _extract_calling_function(source_code, "with_string", 3, Language.PYTHON); result = codeflash_output # 36.2μs -> 30.9μs (17.1% faster)

def test_python_function_with_exception_handling():
    """Test function with try-except blocks."""
    source_code = """def safe_operation():
    try:
        result = 1 / 0
    except ZeroDivisionError:
        result = None
    return result
"""
    codeflash_output = _extract_calling_function(source_code, "safe_operation", 3, Language.PYTHON); result = codeflash_output # 47.4μs -> 41.7μs (13.9% faster)

def test_python_two_functions_same_name_different_scope():
    """Test behavior with functions that have same name (only outer one is found with ast.walk)."""
    source_code = """def duplicate():
    return 1

def wrapper():
    def duplicate():
        return 2
    return duplicate()
"""
    # ast.walk traverses all nodes, so it will find both
    # ref_line = 2 should match the first one
    codeflash_output = _extract_calling_function(source_code, "duplicate", 2, Language.PYTHON); result = codeflash_output # 44.1μs -> 39.6μs (11.4% faster)

def test_javascript_language_with_python_code():
    """Test that JavaScript extraction is attempted for JavaScript language."""
    # When language is not PYTHON, it tries JS extraction
    source_code = "function hello() { return 42; }"
    # This will attempt TreeSitterAnalyzer which might not be fully available
    # but the function should handle exceptions gracefully
    codeflash_output = _extract_calling_function(source_code, "hello", 1, Language.JAVASCRIPT); result = codeflash_output # 54.3μs -> 55.0μs (1.37% slower)

def test_python_function_ref_line_in_string():
    """Test when there's a string containing line number references."""
    source_code = '''def func_with_string():
    code = "def other():\\n    pass"
    return code
'''
    codeflash_output = _extract_calling_function(source_code, "func_with_string", 2, Language.PYTHON); result = codeflash_output # 38.2μs -> 32.0μs (19.5% faster)

def test_python_function_with_type_hints():
    """Test extracting a function with type hints."""
    source_code = """def typed_func(x: int, y: str) -> bool:
    return isinstance(y, str)
"""
    codeflash_output = _extract_calling_function(source_code, "typed_func", 2, Language.PYTHON); result = codeflash_output # 43.6μs -> 38.3μs (13.8% faster)

def test_python_async_function_with_await():
    """Test async function with await statements."""
    source_code = """async def async_op():
    result = await some_coroutine()
    return result
"""
    codeflash_output = _extract_calling_function(source_code, "async_op", 2, Language.PYTHON); result = codeflash_output # 38.0μs -> 32.8μs (15.9% faster)

def test_python_large_source_file_many_functions():
    """Test extracting a function from a large file with many functions."""
    # Create source with 100 functions
    functions = []
    for i in range(100):
        functions.append(f"""def func_{i}():
    return {i}
""")
    source_code = "\n".join(functions)
    
    # Try to extract function 50
    codeflash_output = _extract_calling_function(source_code, "func_50", 202, Language.PYTHON); result = codeflash_output # 927μs -> 505μs (83.4% faster)

def test_python_large_function_with_many_lines():
    """Test extracting a very large function with hundreds of lines."""
    # Create a function with 200 lines
    lines = ["def large_func():"]
    for i in range(200):
        lines.append(f"    x_{i} = {i}")
    lines.append("    return x_199")
    
    source_code = "\n".join(lines)
    
    # ref_line somewhere in the middle
    codeflash_output = _extract_calling_function(source_code, "large_func", 100, Language.PYTHON); result = codeflash_output # 562μs -> 548μs (2.57% faster)

def test_python_deeply_nested_source():
    """Test with deeply nested functions (complexity test)."""
    source_code = """def level_0():
    x = 1
    def level_1():
        y = 2
        def level_2():
            z = 3
            def level_3():
                w = 4
                return w
            return z
        return y
    return x
"""
    codeflash_output = _extract_calling_function(source_code, "level_0", 5, Language.PYTHON); result = codeflash_output # 61.2μs -> 54.8μs (11.7% faster)

def test_python_many_functions_correct_extraction():
    """Test that correct function is extracted among many similar functions."""
    functions = []
    for i in range(50):
        functions.append(f"""def process_{i}(data):
    result = data * {i}
    return result
""")
    source_code = "\n".join(functions)
    
    # Extract process_25
    target_line = sum(len(f.split("\n")) for f in functions[:25]) + 1
    codeflash_output = _extract_calling_function(source_code, "process_25", target_line, Language.PYTHON); result = codeflash_output # 559μs -> 525μs (6.51% faster)

def test_python_large_function_with_complex_logic():
    """Test extracting a large function with complex nested logic."""
    source_code = """def complex_logic():
    for i in range(100):
        if i % 2 == 0:
            for j in range(50):
                try:
                    x = i / j
                except ZeroDivisionError:
                    x = 0
                finally:
                    pass
        else:
            with open('/tmp/file.txt') as f:
                pass
    return None
"""
    codeflash_output = _extract_calling_function(source_code, "complex_logic", 5, Language.PYTHON); result = codeflash_output # 83.9μs -> 80.0μs (4.88% faster)

def test_python_many_decorators_on_single_function():
    """Test function with many decorators."""
    decorators = "\n".join([f"@decorator_{i}" for i in range(20)])
    source_code = f"""{decorators}
def heavily_decorated():
    return 42
"""
    codeflash_output = _extract_calling_function(source_code, "heavily_decorated", 21, Language.PYTHON); result = codeflash_output # 65.2μs -> 57.7μs (12.9% faster)

def test_python_multiple_async_functions():
    """Test extracting specific async function among many."""
    functions = []
    for i in range(50):
        functions.append(f"""async def async_task_{i}():
    await something_{i}()
    return {i}
""")
    source_code = "\n".join(functions)
    
    target_line = sum(len(f.split("\n")) for f in functions[:25]) + 1
    codeflash_output = _extract_calling_function(source_code, "async_task_25", target_line, Language.PYTHON); result = codeflash_output # 475μs -> 433μs (9.59% faster)

def test_python_functions_with_long_parameter_lists():
    """Test functions with very long parameter lists."""
    params = ", ".join([f"param_{i}" for i in range(50)])
    source_code = f"""def many_params({params}):
    return sum([param_0, param_1])
"""
    codeflash_output = _extract_calling_function(source_code, "many_params", 2, Language.PYTHON); result = codeflash_output # 92.3μs -> 86.2μs (7.03% faster)

def test_python_file_with_mixed_content():
    """Test extracting function from file with imports, classes, and functions."""
    source_code = """import os
import sys
from typing import List, Dict

class MyClass:
    def method(self):
        pass

def utility_func():
    return 42

def another_func():
    return utility_func()

x = 10
y = 20
"""
    codeflash_output = _extract_calling_function(source_code, "another_func", 14, Language.PYTHON); result = codeflash_output # 95.5μs -> 62.5μs (52.7% faster)

def test_python_performance_with_very_long_lines():
    """Test performance with functions containing very long lines."""
    long_line = "x = " + " + ".join([str(i) for i in range(500)])
    source_code = f"""def perf_test():
    {long_line}
    return x
"""
    codeflash_output = _extract_calling_function(source_code, "perf_test", 2, Language.PYTHON); result = codeflash_output # 556μs -> 555μs (0.092% faster)

def test_python_extract_from_large_multiline_strings():
    """Test extracting function that contains large multiline strings."""
    source_code = '''def doc_func():
    """
    Very long docstring that spans many lines.
    """ + """
    More docstring content.
    """ * 100
    return None
'''
    codeflash_output = _extract_calling_function(source_code, "doc_func", 3, Language.PYTHON); result = codeflash_output # 39.1μs -> 33.5μs (16.5% 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-pr1256-2026-02-03T00.39.14 and push.