⚡️ Speed up function get_test_info_from_stack by 1,809% in PR #1014 (feat/extract-imported-class-definitions)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1014

If you approve this dependent PR, these changes will be merged into the original PR branch feat/extract-imported-class-definitions.

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

---

📄 1,809% (18.09x) speedup for get_test_info_from_stack in codeflash/verification/codeflash_capture.py

⏱️ Runtime : 29.0 milliseconds → 1.52 milliseconds (best of 62 runs)

📝 Explanation and details

The optimized code achieves an 18x speedup by eliminating two major bottlenecks identified in the line profiler results:

Key Optimizations

1. Replaced inspect.getmodule(frame) with direct dictionary lookup (45.6% → 2.1% of runtime)

• Original: Called inspect.getmodule(frame) which traverses sys.modules and performs expensive introspection
• Optimized: Uses frame.f_globals.get("__name__", None) to directly access the module name from the frame's global namespace
• Impact: Reduces per-hit time from 153,322ns to 444.5ns (~345x faster per call)

2. Optimized path containment check (49.5% → 0.3% of runtime)

• Original: Called Path(filename).resolve().is_relative_to(Path(tests_root)) on every frame, creating Path objects and resolving symlinks repeatedly
• Optimized:
  • Pre-computes tests_root_resolved once using os.path.normcase(os.path.abspath(tests_root))
  • Uses os.path.commonpath() with normalized strings for fast comparison
  • Only executes when conditions match (inside the conditional block)
• Impact: Avoids 543 expensive pathlib operations; when executed, uses string-based operations that are orders of magnitude faster

3. Minor optimization in module name check

• Caches frame.f_globals in a local variable fg to avoid repeated attribute lookups
• Uses conditional expression for safer __name__ retrieval

Why This Matters in Context

Based on the function_references, get_test_info_from_stack() is called on every test invocation within a wrapper that:

• Records timing information for performance testing
• Captures instance state for verification
• Writes results to SQLite

Since this function executes in the hot path of test execution, any overhead directly impacts test suite runtime. The optimization is especially valuable for:

• Large test suites with many test functions (as shown in test_large_scale_repeated_calls: 27.6ms → 1.37ms for 500 calls)
• Tests with deep call stacks where frame walking is more expensive
• CI/CD pipelines where cumulative test time matters

Test Case Performance Patterns

The annotated tests show consistent 8-20x speedups across different scenarios:

• Standard test functions: ~800-900% speedup (e.g., test_basic_test_function_detection: 80.4μs → 8.97μs)
• Repeated calls: ~1900% speedup (500 iterations: 27.6ms → 1.37ms)
• Long paths: ~1350% speedup with complex directory structures
• Edge cases (empty paths, unicode): All maintain similar speedup ratios

The optimization works best when stack walking finds test functions quickly (avoiding the fallback to environment variables), which is the common case in normal test execution.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import inspect
import os
import sys
import types
from pathlib import Path

# imports
import pytest
from codeflash.verification.codeflash_capture import get_test_info_from_stack

# --- Basic Test Cases ---

def test_function_basic_module_level(tmp_path, monkeypatch):
    """
    Test extraction from a simple test function at module level.
    """
    # Simulate a test file in a tests root
    tests_root = tmp_path / "tests"
    tests_root.mkdir()
    test_file = tests_root / "test_sample.py"
    test_file.write_text(
        "def test_foo():\n"
        "    import inspect\n"
        "    from codeflash.verification.codeflash_capture import get_test_info_from_stack\n"
        "    return get_test_info_from_stack('{}')\n"
        "result = test_foo()\n".format(str(tests_root))
    )
    # Dynamically import the test file as a module
    sys.path.insert(0, str(tests_root.parent))
    try:
        import importlib.util
        spec = importlib.util.spec_from_file_location("test_sample", str(test_file))
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
    finally:
        sys.path.pop(0)
        if "test_sample" in sys.modules:
            del sys.modules["test_sample"]

def test_function_basic_class_method(monkeypatch):
    """
    Test extraction from a test method in a class.
    """
    class DummyTestClass:
        

def test_function_no_test_prefix(monkeypatch):
    """
    Should not find test info if not in a test_ function, so returns mostly empty.
    """
    codeflash_output = get_test_info_from_stack(str(Path(__file__).parent)); result = codeflash_output

def test_function_env_fallback(monkeypatch):
    """
    Should fall back to environment variables if no test function in stack.
    """
    monkeypatch.setenv("CODEFLASH_TEST_FUNCTION", "test_env_func")
    monkeypatch.setenv("CODEFLASH_TEST_MODULE", "test_env_module")
    monkeypatch.setenv("CODEFLASH_TEST_CLASS", "TestEnvClass")
    codeflash_output = get_test_info_from_stack(str(Path(__file__).parent)); result = codeflash_output

def test_function_env_partial(monkeypatch):
    """
    Should handle missing class env variable gracefully.
    """
    monkeypatch.setenv("CODEFLASH_TEST_FUNCTION", "test_env_func2")
    monkeypatch.setenv("CODEFLASH_TEST_MODULE", "test_env_module2")
    if "CODEFLASH_TEST_CLASS" in os.environ:
        monkeypatch.delenv("CODEFLASH_TEST_CLASS")
    codeflash_output = get_test_info_from_stack(str(Path(__file__).parent)); result = codeflash_output

def test_function_module_level_stack(monkeypatch, tmp_path):
    """
    Simulate a test module run at module level (function_name == '<module>').
    """
    tests_root = tmp_path / "tests"
    tests_root.mkdir()
    test_file = tests_root / "test_mod.py"
    test_file.write_text(
        "from codeflash.verification.codeflash_capture import get_test_info_from_stack\n"
        "result = get_test_info_from_stack('{}')\n".format(str(tests_root))
    )
    sys.path.insert(0, str(tests_root.parent))
    try:
        import importlib.util
        spec = importlib.util.spec_from_file_location("test_mod", str(test_file))
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
    finally:
        sys.path.pop(0)
        if "test_mod" in sys.modules:
            del sys.modules["test_mod"]


def test_function_class_without_self(monkeypatch):
    """
    Should not detect class name if 'self' is not in locals.
    """
    class Dummy:
        @staticmethod
        
def test_function_class_with_classmethod(monkeypatch):
    """
    Should not detect class name for classmethod (no 'self').
    """
    class Dummy:
        @classmethod
        


def test_function_many_test_functions(monkeypatch):
    """
    Test with many test functions in the stack, should pick the innermost.
    """
    def test_outer():
        
def test_function_large_env(monkeypatch):
    """
    Test with large environment variable values.
    """
    monkeypatch.setenv("CODEFLASH_TEST_FUNCTION", "test_" + "x"*900)
    monkeypatch.setenv("CODEFLASH_TEST_MODULE", "mod_" + "y"*900)
    monkeypatch.setenv("CODEFLASH_TEST_CLASS", "Class" + "z"*900)
    codeflash_output = get_test_info_from_stack(str(Path(__file__).parent)); result = codeflash_output

import os
from pathlib import Path
from unittest.mock import patch

# imports
from codeflash.verification.codeflash_capture import get_test_info_from_stack


# Helper class for testing class-based tests
class TestHelperClass:
    """Helper class to simulate test class methods."""


def non_test_function_helper(tests_root: str):
    """A non-test function that calls get_test_info_from_stack."""
    return get_test_info_from_stack(tests_root)


def nested_call_helper(tests_root: str):
    """Helper that creates nested calls before getting test info."""

    def inner():
        return get_test_info_from_stack(tests_root)

    return inner()


# ============================================================================
# BASIC TEST CASES
# ============================================================================


def test_basic_test_function_detection():
    """Test that get_test_info_from_stack correctly identifies a test function."""
    # Call from within a test function (this function itself)
    tests_root = str(Path(__file__).parent)
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 80.4μs -> 8.97μs (796% faster)


def test_basic_test_class_method_detection():
    """Test that get_test_info_from_stack correctly identifies a test method in a class."""
    tests_root = str(Path(__file__).parent)
    helper = TestHelperClass()
    test_module_name, test_class_name, test_name, line_id = helper.test_method_in_class(tests_root)


def test_basic_nested_call():
    """Test that get_test_info_from_stack works through nested function calls."""
    tests_root = str(Path(__file__).parent)
    test_module_name, test_class_name, test_name, line_id = nested_call_helper(tests_root)


def test_basic_environment_variable_fallback():
    """Test that environment variables are used when stack walking fails."""
    tests_root = str(Path(__file__).parent)

    # Set environment variables
    with patch.dict(
        os.environ,
        {
            "CODEFLASH_TEST_FUNCTION": "test_from_env",
            "CODEFLASH_TEST_MODULE": "module_from_env",
            "CODEFLASH_TEST_CLASS": "ClassFromEnv",
        },
    ):
        # Call from a non-test function to force fallback
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_basic_return_types():
    """Test that return types are correct."""
    tests_root = str(Path(__file__).parent)
    codeflash_output = get_test_info_from_stack(tests_root)
    result = codeflash_output  # 78.1μs -> 8.66μs (802% faster)

    test_module_name, test_class_name, test_name, line_id = result


# ============================================================================
# EDGE TEST CASES
# ============================================================================


def test_edge_empty_tests_root():
    """Test behavior with empty tests_root string."""
    tests_root = ""
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 79.8μs -> 12.9μs (517% faster)


def test_edge_nonexistent_tests_root():
    """Test behavior with non-existent tests_root path."""
    tests_root = "/nonexistent/path/to/tests"
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 84.1μs -> 7.83μs (973% faster)


def test_edge_no_test_prefix():
    """Test that functions without 'test_' prefix are not detected as tests."""
    tests_root = str(Path(__file__).parent)

    # Clear environment variables to ensure no fallback
    with patch.dict(os.environ, {}, clear=True):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_edge_partial_environment_variables():
    """Test behavior when only some environment variables are set."""
    tests_root = str(Path(__file__).parent)

    # Set only test function, not module or class
    with patch.dict(os.environ, {"CODEFLASH_TEST_FUNCTION": "test_partial_env"}, clear=True):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_edge_environment_variable_with_empty_class():
    """Test behavior when CODEFLASH_TEST_CLASS is empty string."""
    tests_root = str(Path(__file__).parent)

    with patch.dict(
        os.environ,
        {"CODEFLASH_TEST_FUNCTION": "test_env", "CODEFLASH_TEST_MODULE": "test_module", "CODEFLASH_TEST_CLASS": ""},
    ):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_edge_class_without_test_prefix():
    """Test that methods in classes without 'test_' prefix are handled correctly."""
    tests_root = str(Path(__file__).parent)
    helper = TestHelperClass()

    # Clear environment variables
    with patch.dict(os.environ, {}, clear=True):
        test_module_name, test_class_name, test_name, line_id = helper.non_test_method(tests_root)


def test_edge_special_characters_in_tests_root():
    """Test behavior with special characters in tests_root path."""
    tests_root = "/path/with spaces/and-dashes/tests_root"
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 87.7μs -> 8.60μs (920% faster)


def test_edge_unicode_in_tests_root():
    """Test behavior with unicode characters in tests_root path."""
    tests_root = "/path/with/unicode/测试/тест"
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 88.1μs -> 8.30μs (960% faster)


def test_edge_relative_tests_root():
    """Test behavior with relative path for tests_root."""
    tests_root = "."
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 78.2μs -> 12.2μs (540% faster)


def test_edge_absolute_tests_root():
    """Test behavior with absolute path for tests_root."""
    tests_root = str(Path(__file__).parent.resolve())
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 67.5μs -> 7.43μs (808% faster)


def test_edge_line_id_is_valid_number():
    """Test that line_id is always a valid string representation of a number."""
    tests_root = str(Path(__file__).parent)
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 73.7μs -> 7.67μs (860% faster)

    # line_id should be convertible to int or empty
    if line_id:
        pass


def test_edge_multiple_test_prefix_in_stack():
    """Test behavior when multiple functions with 'test_' prefix are in the stack."""
    tests_root = str(Path(__file__).parent)

    # This test function calls a helper that also has 'test_' prefix
    test_module_name, test_class_name, test_name, line_id = test_function_helper(tests_root)


def test_edge_environment_override_when_stack_has_test():
    """Test that stack detection takes precedence over environment variables."""
    tests_root = str(Path(__file__).parent)

    # Set environment variables
    with patch.dict(
        os.environ,
        {
            "CODEFLASH_TEST_FUNCTION": "test_from_env",
            "CODEFLASH_TEST_MODULE": "module_from_env",
            "CODEFLASH_TEST_CLASS": "ClassFromEnv",
        },
    ):
        # Call from a test function - stack should take precedence
        test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
            tests_root
        )  # 83.7μs -> 9.49μs (782% faster)


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_large_scale_long_module_name():
    """Test with a very long module name in environment."""
    tests_root = str(Path(__file__).parent)

    long_module = "test_" + "module_" * 100

    with patch.dict(os.environ, {"CODEFLASH_TEST_FUNCTION": "test_func", "CODEFLASH_TEST_MODULE": long_module}):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_large_scale_long_class_name():
    """Test with a very long class name in environment."""
    tests_root = str(Path(__file__).parent)

    long_class = "TestClass" + "Name" * 100

    with patch.dict(os.environ, {"CODEFLASH_TEST_FUNCTION": "test_func", "CODEFLASH_TEST_CLASS": long_class}):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_large_scale_long_tests_root_path():
    """Test with a very long tests_root path."""
    # Create a path with many nested directories
    long_path = "/" + "/".join(["dir" + str(i) for i in range(100)])

    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        long_path
    )  # 145μs -> 10.0μs (1352% faster)


def test_large_scale_repeated_calls():
    """Test performance with many repeated calls."""
    tests_root = str(Path(__file__).parent)

    # Call the function many times
    results = []
    for _ in range(500):
        codeflash_output = get_test_info_from_stack(tests_root)
        result = codeflash_output  # 27.6ms -> 1.37ms (1918% faster)
        results.append(result)
    for result in results:
        test_module_name, test_class_name, test_name, line_id = result


def test_large_scale_many_local_variables():
    """Test behavior when frame has many local variables."""
    tests_root = str(Path(__file__).parent)

    # Create many local variables
    locals_dict = {f"var_{i}": i for i in range(500)}

    # Add them to the local scope
    for key, value in locals_dict.items():
        locals()[key] = value

    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 79.3μs -> 7.88μs (907% faster)


def test_large_scale_many_environment_variables():
    """Test behavior when many environment variables are set."""
    tests_root = str(Path(__file__).parent)

    # Create many environment variables
    env_dict = {f"VAR_{i}": f"value_{i}" for i in range(500)}
    env_dict["CODEFLASH_TEST_FUNCTION"] = "test_from_env"

    with patch.dict(os.environ, env_dict):
        test_module_name, test_class_name, test_name, line_id = non_test_function_helper(tests_root)


def test_large_scale_long_line_numbers():
    """Test that line_id can handle large line numbers."""
    tests_root = str(Path(__file__).parent)

    # We can't actually create a file with millions of lines,
    # but we can verify the function handles the line number correctly
    test_module_name, test_class_name, test_name, line_id = get_test_info_from_stack(
        tests_root
    )  # 77.4μs -> 8.53μs (808% 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-pr1014-2026-01-07T22.27.39 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #1014 by KRRT7 was closed.