⚡️ Speed up function _trim_traceback by 251%

[codeflash-ai]

📄 251% (2.51x) speedup for _trim_traceback in marimo/_messaging/tracebacks.py

⏱️ Runtime : 295 microseconds → 83.8 microseconds (best of 395 runs)

📝 Explanation and details

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
_messaging/test_tracebacks.py::TestTracebacks.test_trim	2.98μs	2.48μs	20.0%✅

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from marimo._messaging.tracebacks import _trim_traceback

# unit tests

# -------------------------
# Basic Test Cases
# -------------------------

def test_trim_traceback_basic_removal():
    """
    Basic: Should remove the first executor frame if present.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    expected = (
        "Traceback (most recent call last):\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.07μs -> 2.41μs (27.7% faster)

def test_trim_traceback_basic_no_removal():
    """
    Basic: Should return unchanged if executor frame not present.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.27μs -> 1.40μs (9.76% slower)

def test_trim_traceback_basic_multiple_frames():
    """
    Basic: Should remove only the first executor frame, keep rest.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        '  File "/marimo/_runtime/executor.py", line 99, in something_else\n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    expected = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 99, in something_else\n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.13μs -> 2.37μs (32.2% faster)

# -------------------------
# Edge Test Cases
# -------------------------

def test_trim_traceback_empty_string():
    """
    Edge: Should return empty string unchanged.
    """
    tb = ""
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 623ns -> 542ns (14.9% faster)

def test_trim_traceback_single_line():
    """
    Edge: Should return single-line string unchanged.
    """
    tb = "ValueError: error"
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 645ns -> 566ns (14.0% faster)

def test_trim_traceback_two_lines():
    """
    Edge: Should return two-line string unchanged.
    """
    tb = "Traceback (most recent call last):\nValueError: error"
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 748ns -> 1.18μs (36.8% slower)

def test_trim_traceback_executor_frame_not_at_second_line():
    """
    Edge: Should not trim if executor frame is not at index 1.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.34μs -> 1.45μs (7.58% slower)

def test_trim_traceback_executor_frame_wrong_suffix():
    """
    Edge: Should not trim if executor frame does not end with correct suffix.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in not_execute_cell\n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.57μs -> 1.66μs (5.48% slower)

def test_trim_traceback_executor_frame_wrong_prefix():
    """
    Edge: Should not trim if executor frame does not contain correct prefix.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/other.py", line 42, in execute_cell\n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.30μs -> 1.42μs (8.80% slower)

def test_trim_traceback_no_file_lines_after_executor():
    """
    Edge: Should not trim if no '  File ' line after executor frame.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.37μs -> 1.91μs (24.5% faster)

def test_trim_traceback_nonstandard_format():
    """
    Edge: Should return unchanged for nonstandard traceback format.
    """
    tb = (
        "Some random error string\n"
        "Not a traceback"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 761ns -> 782ns (2.69% slower)

def test_trim_traceback_executor_frame_with_extra_whitespace():
    """
    Edge: Should not trim if executor frame line has extra whitespace.
    """
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell   \n'
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.59μs -> 1.61μs (0.809% slower)

def test_trim_traceback_executor_frame_with_different_quote():
    """
    Edge: Should not trim if executor frame uses single quotes instead of double quotes.
    """
    tb = (
        "Traceback (most recent call last):\n"
        "  File '/marimo/_runtime/executor.py', line 42, in execute_cell\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.27μs -> 1.41μs (10.1% slower)

# -------------------------
# Large Scale Test Cases
# -------------------------

def test_trim_traceback_large_traceback():
    """
    Large Scale: Should trim correctly with large number of frames.
    """
    # Build a large traceback with 500 frames after executor
    user_frames = [
        f'  File "/path/to/user_code{i}.py", line {i}, in <module>'
        for i in range(1, 501)
    ]
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        + "\n".join(user_frames) + "\n"
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    expected = (
        "Traceback (most recent call last):\n"
        + "\n".join(user_frames) + "\n"
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 23.5μs -> 4.24μs (455% faster)

def test_trim_traceback_large_traceback_no_executor():
    """
    Large Scale: Should return unchanged for large traceback with no executor frame.
    """
    user_frames = [
        f'  File "/path/to/user_code{i}.py", line {i}, in <module>'
        for i in range(1, 501)
    ]
    tb = (
        "Traceback (most recent call last):\n"
        + "\n".join(user_frames) + "\n"
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 14.6μs -> 1.43μs (922% faster)

def test_trim_traceback_large_traceback_executor_not_at_second_line():
    """
    Large Scale: Should not trim if executor frame is not at index 1.
    """
    user_frames = [
        f'  File "/path/to/user_code{i}.py", line {i}, in <module>'
        for i in range(1, 500)
    ]
    tb = (
        "Traceback (most recent call last):\n"
        + "\n".join(user_frames) + "\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 14.0μs -> 1.42μs (888% faster)

def test_trim_traceback_large_traceback_multiple_executor_frames():
    """
    Large Scale: Should only remove the first executor frame at index 1.
    """
    user_frames = [
        f'  File "/marimo/_runtime/executor.py", line {i}, in execute_cell'
        for i in range(43, 53)
    ]
    tb = (
        "Traceback (most recent call last):\n"
        '  File "/marimo/_runtime/executor.py", line 42, in execute_cell\n'
        + "\n".join(user_frames) + "\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    expected = (
        "Traceback (most recent call last):\n"
        + "\n".join(user_frames) + "\n"
        '  File "/path/to/user_code.py", line 10, in <module>\n'
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.23μs -> 2.53μs (27.5% faster)

def test_trim_traceback_large_traceback_all_executor_frames():
    """
    Large Scale: If all frames after header are executor frames, should keep all except the first.
    """
    executor_frames = [
        f'  File "/marimo/_runtime/executor.py", line {i}, in execute_cell'
        for i in range(42, 1042)
    ]
    tb = (
        "Traceback (most recent call last):\n"
        + "\n".join(executor_frames) + "\n"
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    expected = (
        "Traceback (most recent call last):\n"
        + "\n".join(executor_frames[1:]) + "\n"
        "    raise ValueError('error')\n"
        "ValueError: error"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 77.2μs -> 6.03μs (1181% 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  # used for our unit tests
from marimo._messaging.tracebacks import _trim_traceback

# unit tests

# ------------------- Basic Test Cases -------------------

def test_no_trim_needed_non_matching_traceback():
    # Traceback does not match the pattern, should not be trimmed
    tb = "Traceback (most recent call last):\n  File \"otherfile.py\", line 10, in <module>\nValueError: error"
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.27μs -> 1.67μs (24.0% slower)

def test_trim_matching_traceback():
    # Traceback matches the pattern, should be trimmed
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 123, in execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    expected = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.19μs -> 2.47μs (29.1% faster)

def test_trim_with_multiple_files():
    # Multiple "File" lines, should trim only the first executor line
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "  File \"user_code.py\", line 2, in foo\n"
        "  File \"lib.py\", line 3, in bar\n"
        "    1/0\n"
        "ZeroDivisionError: division by zero"
    )
    expected = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 2, in foo\n"
        "  File \"lib.py\", line 3, in bar\n"
        "    1/0\n"
        "ZeroDivisionError: division by zero"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.90μs -> 2.35μs (23.3% faster)

def test_no_file_line_after_executor():
    # No "File" line after executor, should not trim
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "Some other line\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.24μs -> 1.81μs (23.9% faster)

def test_traceback_with_only_executor():
    # Only executor line after header, should not trim
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.11μs -> 1.85μs (14.3% faster)

# ------------------- Edge Test Cases -------------------

def test_empty_traceback():
    # Empty string input
    tb = ""
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 628ns -> 523ns (20.1% faster)

def test_single_line_traceback():
    # Only one line, not a traceback
    tb = "ValueError: bad"
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 655ns -> 515ns (27.2% faster)

def test_two_line_traceback():
    # Only two lines, not enough to match the pattern
    tb = "Traceback (most recent call last):\nValueError: bad"
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 767ns -> 1.13μs (32.4% slower)

def test_executor_line_not_at_second_position():
    # Executor line is not the second line, should not trim
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.30μs -> 1.51μs (13.4% slower)

def test_executor_line_wrong_suffix():
    # Executor line does not end with ', in execute_cell'
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in not_execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.53μs -> 1.70μs (10.4% slower)

def test_executor_line_wrong_prefix():
    # Executor line does not contain the expected prefix
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/other/path/executor.py\", line 99, in execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 1.23μs -> 1.50μs (18.0% slower)

def test_executor_line_with_extra_spaces():
    # Executor line with extra spaces, should not trim
    tb = (
        "Traceback (most recent call last):\n"
        "   File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.24μs -> 2.42μs (33.7% faster)

def test_file_line_not_indented():
    # "File" line after executor is not properly indented
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    # Should not trim, as file line does not start with "  File "
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.31μs -> 1.97μs (17.3% faster)

def test_file_line_with_tabs():
    # "File" line after executor uses tabs instead of spaces
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "\tFile \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    # Should not trim, as file line does not start with "  File "
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.22μs -> 1.96μs (13.4% faster)

def test_traceback_with_intervening_non_file_lines():
    # There are non-"File" lines between executor and next "File" line
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "Some log message\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    expected = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.12μs -> 2.38μs (30.8% faster)

def test_traceback_with_multiple_executors():
    # Multiple executor entries, only the first should be trimmed
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "  File \"/marimo/_runtime/executor.py\", line 101, in execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    expected = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bad')\n"
        "ValueError: bad"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 2.87μs -> 2.28μs (26.2% faster)

def test_traceback_with_unicode_characters():
    # Unicode characters in traceback
    tb = (
        "Traceback (most recent call last):\n"
        "  File \"/marimo/_runtime/executor.py\", line 99, in execute_cell\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bäd')\n"
        "ValueError: bäd"
    )
    expected = (
        "Traceback (most recent call last):\n"
        "  File \"user_code.py\", line 5, in <module>\n"
        "    raise ValueError('bäd')\n"
        "ValueError: bäd"
    )
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 3.45μs -> 3.01μs (14.4% faster)

# ------------------- Large Scale Test Cases -------------------

def test_large_traceback_trim():
    # Large traceback with executor at start, should trim only the first executor entry
    user_file_lines = [
        f"  File \"user_code.py\", line {i}, in func{i}" for i in range(1, 500)
    ]
    tb_lines = [
        "Traceback (most recent call last):",
        "  File \"/marimo/_runtime/executor.py\", line 42, in execute_cell",
    ] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ]
    tb = "\n".join(tb_lines)
    expected = "\n".join([tb_lines[0]] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ])
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 20.7μs -> 3.88μs (434% faster)

def test_large_traceback_no_trim():
    # Large traceback without executor, should not be trimmed
    user_file_lines = [
        f"  File \"user_code.py\", line {i}, in func{i}" for i in range(1, 500)
    ]
    tb_lines = [
        "Traceback (most recent call last):",
    ] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ]
    tb = "\n".join(tb_lines)
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 12.4μs -> 1.42μs (771% faster)

def test_large_traceback_multiple_executors():
    # Large traceback with multiple executor entries, only the first should be trimmed
    user_file_lines = [
        f"  File \"user_code.py\", line {i}, in func{i}" for i in range(1, 500)
    ]
    tb_lines = [
        "Traceback (most recent call last):",
        "  File \"/marimo/_runtime/executor.py\", line 42, in execute_cell",
        "  File \"/marimo/_runtime/executor.py\", line 43, in execute_cell",
    ] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ]
    expected = "\n".join([tb_lines[0]] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ])
    tb = "\n".join(tb_lines)
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 18.8μs -> 3.43μs (448% faster)

def test_large_traceback_with_intervening_lines():
    # Large traceback with non-"File" lines between executor and file lines
    user_file_lines = [
        f"  File \"user_code.py\", line {i}, in func{i}" for i in range(1, 500)
    ]
    tb_lines = [
        "Traceback (most recent call last):",
        "  File \"/marimo/_runtime/executor.py\", line 42, in execute_cell",
        "Some debug info",
    ] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ]
    expected = "\n".join([tb_lines[0]] + user_file_lines + [
        "    raise Exception('big error')",
        "Exception: big error"
    ])
    tb = "\n".join(tb_lines)
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 18.9μs -> 3.34μs (465% faster)

def test_large_traceback_with_no_file_lines():
    # Large traceback with executor, but no "File" lines after
    tb_lines = [
        "Traceback (most recent call last):",
        "  File \"/marimo/_runtime/executor.py\", line 42, in execute_cell",
    ] + ["Some debug info" for _ in range(500)] + [
        "    raise Exception('big error')",
        "Exception: big error"
    ]
    tb = "\n".join(tb_lines)
    codeflash_output = _trim_traceback(tb); result = codeflash_output # 32.4μs -> 3.04μs (967% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from marimo._messaging.tracebacks import _trim_traceback

def test__trim_traceback():
    _trim_traceback('\n\n')

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_hg3s6k0k/tmpox47i8lk/test_concolic_coverage.py::test__trim_traceback	961ns	822ns	16.9%✅

To edit these changes git checkout codeflash/optimize-_trim_traceback-mhb8nazr and push.