⚡️ Speed up function extract_test_context_from_frame by 491% in PR #687 (granular-async-instrumentation)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #687

If you approve this dependent PR, these changes will be merged into the original PR branch granular-async-instrumentation.

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

---

📄 491% (4.91x) speedup for extract_test_context_from_frame in codeflash/code_utils/codeflash_wrap_decorator.py

⏱️ Runtime : 9.55 milliseconds → 1.62 milliseconds (best of 128 runs)

📝 Explanation and details

The optimization replaces expensive Path object creation and method calls with direct string manipulation operations, delivering a 491% speedup.

Key optimizations:

1. Eliminated Path object overhead: Replaced Path(filename).stem.startswith("test_") with filename.rpartition('/')[-1].rpartition('\\')[-1].rpartition('.')[0].startswith("test_") - avoiding Path instantiation entirely.

2. Optimized path parts extraction: Replaced Path(filename).parts with filename.replace('\\', '/').split('/') - using simple string operations instead of Path parsing.

Performance impact analysis:

• Original profiler shows lines 25-26 (Path operations) consumed 86.3% of total runtime (44.7% + 41.6%)
• Optimized version reduces these same operations to just 25.4% of runtime (15% + 10.4%)
• The string manipulation operations are ~6x faster per call than Path object creation

Test case benefits:

• Large-scale tests see the biggest gains (516% faster for 900-frame stack, 505% faster for 950-frame chain) because the Path overhead multiplies with stack depth
• Edge cases with complex paths benefit significantly (182-206% faster for subdirectory and pytest frame tests)
• Basic tests show minimal overhead since Path operations weren't the bottleneck in shallow stacks

The optimization maintains identical behavior while eliminating the most expensive operations identified in the profiling data - Path object instantiation and method calls that occurred once per stack frame.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import inspect
import sys
import types
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.codeflash_wrap_decorator import \
    extract_test_context_from_frame

# unit tests

# --- Basic Test Cases ---


def test_basic_class_method_detection(monkeypatch):
    """Test detection of a test method inside a class named 'TestBar'."""
    class TestBar:
        def test_bar(self):
            return extract_test_context_from_frame()
    instance = TestBar()
    mod_name = "test_class_module"
    monkeypatch.setitem(instance.test_bar.__globals__, "__name__", mod_name)
    result = instance.test_bar()









def test_multiple_potential_tests(monkeypatch):
    """Test that the first test function is returned when multiple are found."""
    class TestOne:
        def test_one(self):
            return extract_test_context_from_frame()
    class TestTwo:
        def test_two(self):
            return extract_test_context_from_frame()
    instance1 = TestOne()
    instance2 = TestTwo()
    mod_name = "test_multi_module"
    monkeypatch.setitem(instance1.test_one.__globals__, "__name__", mod_name)
    monkeypatch.setitem(instance2.test_two.__globals__, "__name__", mod_name)
    # Call both, only first should be returned
    result1 = instance1.test_one()
    result2 = instance2.test_two()

# --- Large Scale Test Cases ---





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

import inspect
import sys
import types
from pathlib import Path

# imports
import pytest  # used for our unit tests
from codeflash.code_utils.codeflash_wrap_decorator import \
    extract_test_context_from_frame

# ----------------------------------------
# UNIT TESTS FOR extract_test_context_from_frame
# ----------------------------------------

# Helper to simulate a stack frame chain
def make_fake_stack(frames):
    """Creates a chain of fake frames for testing."""
    class FakeFrame:
        def __init__(self, f_code, f_globals, f_locals, f_back, f_filename):
            self.f_code = f_code
            self.f_globals = f_globals
            self.f_locals = f_locals
            self.f_back = f_back
            self.f_code.co_name = f_code.co_name
            self.f_code.co_filename = f_filename
    prev = None
    for frame in reversed(frames):
        f_code = types.SimpleNamespace(co_name=frame['co_name'], co_filename=frame['co_filename'])
        f_globals = frame.get('f_globals', {})
        f_locals = frame.get('f_locals', {})
        f_filename = frame['co_filename']
        curr = FakeFrame(f_code, f_globals, f_locals, prev, f_filename)
        prev = curr
    return prev

# -------- BASIC TEST CASES --------
def test_basic_test_function(monkeypatch):
    """Should detect a basic test function named test_foo."""
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_module.py'},
        {'co_name': 'test_foo', 'co_filename': '/project/test_module.py', 'f_globals': {'__name__': 'test_module'}}
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 1.88μs -> 1.93μs (2.53% slower)

def test_basic_test_class_method(monkeypatch):
    """Should detect a test method inside a class named TestBar."""
    class TestBar:
        pass
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_module.py'},
        {'co_name': 'test_bar', 'co_filename': '/project/test_module.py', 'f_globals': {'__name__': 'test_module'}, 'f_locals': {'self': TestBar()}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 2.48μs -> 2.49μs (0.762% slower)

def test_basic_test_module_name(monkeypatch):
    """Should detect a test context from a module named test_mod."""
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_mod.py'},
        {'co_name': 'some_helper', 'co_filename': '/project/test_mod.py', 'f_globals': {'__name__': 'test_mod'}}
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 3.30μs -> 3.23μs (2.17% faster)

# -------- EDGE TEST CASES --------
def test_edge_no_test_function(monkeypatch):
    """Should raise RuntimeError if no test context is found."""
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/module.py'},
        {'co_name': 'helper', 'co_filename': '/project/module.py', 'f_globals': {'__name__': 'module'}}
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    with pytest.raises(RuntimeError):
        extract_test_context_from_frame() # 21.0μs -> 9.00μs (133% faster)

def test_edge_test_in_subdirectory(monkeypatch):
    """Should detect test context if filename is in a 'test' subdirectory."""
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/tests/test_mod.py'},
        {'co_name': 'helper', 'co_filename': '/project/tests/test_mod.py', 'f_globals': {'__name__': 'tests.test_mod'}}
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 11.9μs -> 4.22μs (182% faster)

def test_edge_class_name_with_test(monkeypatch):
    """Should detect test context if class name contains 'test'."""
    class IntegrationTest:
        pass
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/integration_test.py'},
        {'co_name': 'runTest', 'co_filename': '/project/integration_test.py', 'f_globals': {'__name__': 'integration_test'}, 'f_locals': {'self': IntegrationTest()}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 17.2μs -> 5.81μs (197% faster)

def test_edge_pytest_in_file(monkeypatch):
    """Should detect test context if frame is from pytest file."""
    class TestPy:
        pass
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_py.py'},
        {'co_name': 'run', 'co_filename': '/usr/lib/python3.10/site-packages/pytest.py', 'f_globals': {'__name__': 'pytest', '__file__': '/usr/lib/python3.10/site-packages/pytest.py'}, 'f_locals': {'self': TestPy()}},
        {'co_name': 'test_func', 'co_filename': '/project/test_py.py', 'f_globals': {'__name__': 'test_py'}, 'f_locals': {'self': TestPy()}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 18.8μs -> 6.13μs (206% faster)

def test_edge_test_method_name(monkeypatch):
    """Should detect test context if _testMethodName is present in self."""
    class TestEdge:
        def __init__(self):
            self._testMethodName = "test_edge_case"
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_edge.py'},
        {'co_name': 'runTest', 'co_filename': '/project/test_edge.py', 'f_globals': {'__name__': 'test_edge'}, 'f_locals': {'self': TestEdge()}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 4.21μs -> 4.37μs (3.69% slower)

def test_edge_filename_startswith_test(monkeypatch):
    """Should detect test context if filename starts with test_."""
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_special.py'},
        {'co_name': 'special_func', 'co_filename': '/project/test_special.py', 'f_globals': {'__name__': 'test_special'}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 3.12μs -> 3.10μs (0.646% faster)

# -------- LARGE SCALE TEST CASES --------
def test_large_scale_many_frames(monkeypatch):
    """Should handle a large stack with many intermediate frames and find the test context."""
    # 900 non-test frames, then a test frame
    frames = []
    for i in range(900):
        frames.append({'co_name': f'helper_{i}', 'co_filename': '/project/module.py', 'f_globals': {'__name__': 'module'}})
    # The test frame
    frames.append({'co_name': 'test_massive', 'co_filename': '/project/test_massive.py', 'f_globals': {'__name__': 'test_massive'}})
    # The extract_test_context_from_frame frame
    frames.insert(0, {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_massive.py'})
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 4.60ms -> 747μs (516% faster)

def test_large_scale_many_potential_tests(monkeypatch):
    """Should return the first test context if multiple potential test contexts are found."""
    # Multiple frames that could be tests, only one starts with test_
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_multi.py'},
        {'co_name': 'helper_func', 'co_filename': '/project/test_multi.py', 'f_globals': {'__name__': 'test_multi'}},
        {'co_name': 'test_true', 'co_filename': '/project/test_multi.py', 'f_globals': {'__name__': 'test_multi'}},
        {'co_name': 'test_false', 'co_filename': '/project/test_multi.py', 'f_globals': {'__name__': 'test_multi'}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 3.33μs -> 3.48μs (4.29% slower)

def test_large_scale_long_chain_with_class(monkeypatch):
    """Should handle a long chain of frames with a test method in a class at the end."""
    class TestLong:
        pass
    frames = []
    for i in range(950):
        frames.append({'co_name': f'func_{i}', 'co_filename': '/project/long_module.py', 'f_globals': {'__name__': 'long_module'}})
    frames.append({'co_name': 'test_long_chain', 'co_filename': '/project/test_long_chain.py', 'f_globals': {'__name__': 'test_long_chain'}, 'f_locals': {'self': TestLong()}})
    frames.insert(0, {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_long_chain.py'})
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 4.84ms -> 800μs (505% faster)

def test_large_scale_multiple_test_classes(monkeypatch):
    """Should handle multiple test classes and methods in the stack."""
    class TestA: pass
    class TestB: pass
    frames = [
        {'co_name': 'extract_test_context_from_frame', 'co_filename': '/project/test_multi_class.py'},
        {'co_name': 'helper', 'co_filename': '/project/test_multi_class.py', 'f_globals': {'__name__': 'test_multi_class'}},
        {'co_name': 'test_a', 'co_filename': '/project/test_multi_class.py', 'f_globals': {'__name__': 'test_multi_class'}, 'f_locals': {'self': TestA()}},
        {'co_name': 'test_b', 'co_filename': '/project/test_multi_class.py', 'f_globals': {'__name__': 'test_multi_class'}, 'f_locals': {'self': TestB()}},
    ]
    fake_stack = make_fake_stack(frames)
    inspect._orig_currentframe = inspect.currentframe
    monkeypatch.setattr(inspect, "currentframe", lambda: fake_stack)
    codeflash_output = extract_test_context_from_frame(); result = codeflash_output # 4.04μs -> 4.20μs (3.81% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from codeflash.code_utils.codeflash_wrap_decorator import extract_test_context_from_frame
import pytest

def test_extract_test_context_from_frame():
    with pytest.raises(RuntimeError, match='No\\ test\\ function\\ found\\ in\\ call\\ stack'):
        extract_test_context_from_frame()

To edit these changes git checkout codeflash/optimize-pr687-2025-08-29T22.23.28 and push.

[misrasaurabh1]

it does like repeated work in creating the Path objects everytime is expensive. @KRRT7 why don't we create the path object at the top just once? Can you also see if that candidate was generated during the optimization attempt?

[KRRT7]

it does like repeated work in creating the Path objects everytime is expensive. @KRRT7 why don't we create the path object at the top just once? Can you also see if that candidate was generated during the optimization attempt?

yep, I'm sad that codeflash didn't think of that though

[KRRT7]

applied opt upstream