⚡️ Speed up function create_trace_replay_test_code by 32% in PR #363 (part-1-windows-fixes)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #363

If you approve this dependent PR, these changes will be merged into the original PR branch part-1-windows-fixes.

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

---

📄 32% (0.32x) speedup for create_trace_replay_test_code in codeflash/benchmarking/replay_test.py

⏱️ Runtime : 23.4 milliseconds → 17.7 milliseconds (best of 250 runs)

📝 Explanation and details

Here's a runtime-optimized rewrite based directly on the profile. The main high-cost issues are.

1. Repeated module/function alias generation (especially in two function_imports loops) — avoid recomputing!
2. Many Path(file_path).as_posix() calls — precompute or cache since file_path is not mutated.
3. Constant string formatting and dedent/indent in loops — minimize usage.
4. Code structure — restrain from repeated lookups and temporary allocations in large inner loops.

Below is the optimized code (all results/behavior unchanged).

Key improvements.

• All function/class/alias/file_path strings are now generated once up-front and referenced by mapping, not recomputed every iteration.
• No .split/.join calls or Path() constructions in inner loops.
• No textwrap.dedent/indent in inner loop. Uses fast string join with one indentation pass.
• Eliminated duplicate lookups in functions_data.
• Minimized unnecessary set/list and string allocations.

This should yield a significant performance boost (especially at scale). Output/behavior is identical; semantic minimization confirmed.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import textwrap
from pathlib import Path
from typing import Any

# imports
import pytest
from codeflash.benchmarking.replay_test import create_trace_replay_test_code


# Helper mock for function_properties
class MockFunctionProperties:
    def __init__(self, is_classmethod=False, is_staticmethod=False):
        self.is_classmethod = is_classmethod
        self.is_staticmethod = is_staticmethod

# ---------------------- UNIT TESTS ----------------------

# 1. BASIC TEST CASES

def test_single_function_pytest():
    # Test basic generation for a single function, pytest style
    trace_file = "mytrace.sqlite"
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/path/to/foo/bar.py",
            "benchmark_function_name": "baz",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="pytest", max_run_count=10); code = codeflash_output # 99.7μs -> 26.1μs (282% faster)

def test_single_function_unittest():
    # Test basic generation for a single function, unittest style
    trace_file = "trace.db"
    functions_data = [
        {
            "module_name": "mod",
            "function_name": "func",
            "file_path": "mod.py",
            "benchmark_function_name": "func",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="unittest", max_run_count=5); code = codeflash_output # 95.3μs -> 21.9μs (335% faster)

def test_multiple_functions_sorted_and_no_init():
    # Test that multiple functions are sorted and __init__ is excluded from functions list
    trace_file = "t.sqlite"
    functions_data = [
        {
            "module_name": "pkg.a",
            "function_name": "alpha",
            "file_path": "pkg/a.py",
            "benchmark_function_name": "alpha",
            "function_properties": MockFunctionProperties(),
        },
        {
            "module_name": "pkg.b",
            "function_name": "beta",
            "file_path": "pkg/b.py",
            "benchmark_function_name": "beta",
            "function_properties": MockFunctionProperties(),
        },
        {
            "module_name": "pkg.c",
            "function_name": "__init__",
            "file_path": "pkg/c.py",
            "benchmark_function_name": "__init__",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 115μs -> 41.3μs (179% faster)

def test_function_and_class_method():
    # Test function and a class method (non-static, non-classmethod)
    trace_file = "db.sqlite"
    functions_data = [
        {
            "module_name": "mod",
            "function_name": "func",
            "file_path": "mod.py",
            "benchmark_function_name": "func",
            "function_properties": MockFunctionProperties(),
        },
        {
            "module_name": "mod",
            "function_name": "meth",
            "class_name": "MyClass",
            "file_path": "mod.py",
            "benchmark_function_name": "meth",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 108μs -> 31.7μs (243% faster)

def test_classmethod_and_staticmethod():
    # Test classmethod and staticmethod handling
    trace_file = "trace.db"
    functions_data = [
        {
            "module_name": "m",
            "function_name": "cm",
            "class_name": "C",
            "file_path": "m.py",
            "benchmark_function_name": "cm",
            "function_properties": MockFunctionProperties(is_classmethod=True),
        },
        {
            "module_name": "m",
            "function_name": "sm",
            "class_name": "C",
            "file_path": "m.py",
            "benchmark_function_name": "sm",
            "function_properties": MockFunctionProperties(is_staticmethod=True),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 108μs -> 30.9μs (250% faster)

# 2. EDGE TEST CASES

def test_empty_functions_data():
    # Test with empty functions_data list
    trace_file = "db.sqlite"
    functions_data = []
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 76.2μs -> 4.62μs (1550% faster)

def test_function_name_with_underscore_and_module_with_dot():
    # Test function name with underscores and module with multiple dots
    trace_file = "trace.sqlite"
    functions_data = [
        {
            "module_name": "package.sub.mod",
            "function_name": "my_func_1",
            "file_path": "package/sub/mod.py",
            "benchmark_function_name": "my_func_1",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 97.6μs -> 23.1μs (323% faster)

def test_function_with_init_name_in_class():
    # Test that __init__ in a class is handled but not in functions list
    trace_file = "db.sqlite"
    functions_data = [
        {
            "module_name": "mod",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "mod.py",
            "benchmark_function_name": "__init__",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 96.2μs -> 23.8μs (304% faster)

def test_non_string_trace_file_path():
    # Test that trace_file can be a Path object and is stringified
    trace_file = Path("abc/trace.sqlite")
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "file_path": "foo.py",
            "benchmark_function_name": "bar",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 92.1μs -> 20.5μs (350% faster)

def test_function_with_missing_optional_fields():
    # Test that missing class_name is handled gracefully
    trace_file = "trace.sqlite"
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "file_path": "foo.py",
            "benchmark_function_name": "bar",
            "function_properties": MockFunctionProperties(),
            # class_name missing
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 93.1μs -> 19.4μs (380% faster)

def test_invalid_test_framework():
    # Test that invalid test_framework raises AssertionError
    trace_file = "trace.sqlite"
    functions_data = []
    with pytest.raises(AssertionError):
        create_trace_replay_test_code(trace_file, functions_data, test_framework="nose")

def test_method_with_classname_containing_underscore():
    # Test that class_name with underscores is handled in alias
    trace_file = "trace.sqlite"
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": "My_Class",
            "file_path": "foo.py",
            "benchmark_function_name": "bar",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 97.2μs -> 25.8μs (277% faster)

def test_method_with_file_path_backslashes():
    # Test that file_path with backslashes is normalized to posix
    trace_file = "trace.sqlite"
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": "C",
            "file_path": r"foo\bar.py",
            "benchmark_function_name": "bar",
            "function_properties": MockFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 97.8μs -> 24.4μs (301% faster)

# 3. LARGE SCALE TEST CASES

def test_many_functions_large_scale():
    # Test with a large number of functions (but < 1000)
    trace_file = "trace.sqlite"
    n = 100
    functions_data = [
        {
            "module_name": f"pkg.m{i}",
            "function_name": f"f{i}",
            "file_path": f"pkg/m{i}.py",
            "benchmark_function_name": f"f{i}",
            "function_properties": MockFunctionProperties(),
        }
        for i in range(n)
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 788μs -> 673μs (17.0% faster)
    # All imports present
    for i in range(n):
        pass
    # Functions list should include all function names, sorted
    sorted_names = sorted([f"f{i}" for i in range(n)])

def test_many_methods_large_scale():
    # Test with a large number of methods (but < 1000)
    trace_file = "trace.sqlite"
    n = 50
    functions_data = [
        {
            "module_name": f"pkg.m{i}",
            "function_name": f"meth{i}",
            "class_name": f"C{i}",
            "file_path": f"pkg/m{i}.py",
            "benchmark_function_name": f"meth{i}",
            "function_properties": MockFunctionProperties(),
        }
        for i in range(n)
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output # 592μs -> 490μs (20.8% faster)
    # All class imports and test methods present
    for i in range(n):
        pass

def test_large_scale_mixed_functions_and_methods():
    # Test with a mix of functions and methods, large scale
    trace_file = "trace.sqlite"
    n = 40
    functions_data = []
    for i in range(n):
        # Even: function, Odd: method
        if i % 2 == 0:
            functions_data.append({
                "module_name": f"pkg.m{i}",
                "function_name": f"f{i}",
                "file_path": f"pkg/m{i}.py",
                "benchmark_function_name": f"f{i}",
                "function_properties": MockFunctionProperties(),
            })
        else:
            functions_data.append({
                "module_name": f"pkg.m{i}",
                "function_name": f"meth{i}",
                "class_name": f"C{i}",
                "file_path": f"pkg/m{i}.py",
                "benchmark_function_name": f"meth{i}",
                "function_properties": MockFunctionProperties(),
            })
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); code = codeflash_output
    for i in range(n):
        if i % 2 == 0:
            pass
        else:
            pass

def test_large_scale_max_run_count():
    # Test that large max_run_count is respected in all test bodies
    trace_file = "trace.sqlite"
    n = 20
    max_run_count = 999
    functions_data = [
        {
            "module_name": f"pkg.m{i}",
            "function_name": f"f{i}",
            "file_path": f"pkg/m{i}.py",
            "benchmark_function_name": f"f{i}",
            "function_properties": MockFunctionProperties(),
        }
        for i in range(n)
    ]
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, max_run_count=max_run_count); code = codeflash_output # 235μs -> 151μs (55.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import textwrap
from pathlib import Path
from typing import Any

# imports
import pytest  # used for our unit tests
from codeflash.benchmarking.replay_test import create_trace_replay_test_code


# Helper class for function_properties mock
class DummyFunctionProperties:
    def __init__(self, is_classmethod=False, is_staticmethod=False):
        self.is_classmethod = is_classmethod
        self.is_staticmethod = is_staticmethod

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

def test_basic_single_function_pytest():
    # Test a single top-level function, pytest style
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": "",
            "file_path": "/path/to/foo/bar.py",
            "benchmark_function_name": "baz",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=10,
    ); code = codeflash_output # 99.8μs -> 24.6μs (305% faster)

def test_basic_single_function_unittest():
    # Test a single top-level function, unittest style
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": "",
            "file_path": "/path/to/foo/bar.py",
            "benchmark_function_name": "baz",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="unittest",
        max_run_count=5,
    ); code = codeflash_output # 97.2μs -> 24.0μs (305% faster)

def test_basic_multiple_functions():
    # Test multiple top-level functions from different modules
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "alpha",
            "class_name": "",
            "file_path": "/foo.py",
            "benchmark_function_name": "alpha",
            "function_properties": DummyFunctionProperties(),
        },
        {
            "module_name": "bar",
            "function_name": "beta",
            "class_name": "",
            "file_path": "/bar.py",
            "benchmark_function_name": "beta",
            "function_properties": DummyFunctionProperties(),
        },
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 106μs -> 31.8μs (236% faster)

def test_basic_class_method_and_static_method():
    # Test class method and static method generation
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "classy",
            "class_name": "MyClass",
            "file_path": "/foo.py",
            "benchmark_function_name": "classy",
            "function_properties": DummyFunctionProperties(is_classmethod=True),
        },
        {
            "module_name": "foo",
            "function_name": "staticy",
            "class_name": "MyClass",
            "file_path": "/foo.py",
            "benchmark_function_name": "staticy",
            "function_properties": DummyFunctionProperties(is_staticmethod=True),
        },
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=3,
    ); code = codeflash_output # 109μs -> 31.5μs (246% faster)

def test_basic_init_method():
    # Test __init__ method handling
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/foo.py",
            "benchmark_function_name": "__init__",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=2,
    ); code = codeflash_output # 96.7μs -> 24.0μs (303% faster)

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

def test_empty_functions_data():
    # Test with empty functions_data
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=[],
        test_framework="pytest",
        max_run_count=10,
    ); code = codeflash_output # 75.3μs -> 4.70μs (1502% faster)

def test_function_name_with_underscore_and_dots():
    # Function names and module names with underscores and dots
    functions_data = [
        {
            "module_name": "foo.bar_baz",
            "function_name": "qux_quux",
            "class_name": "",
            "file_path": "/foo/bar_baz.py",
            "benchmark_function_name": "qux_quux",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=4,
    ); code = codeflash_output # 97.0μs -> 23.6μs (310% faster)

def test_path_conversion_forward_slash():
    # File paths should be converted to posix (forward slash)
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": "",
            "file_path": "foo\\bar.py",
            "benchmark_function_name": "bar",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=7,
    ); code = codeflash_output # 95.2μs -> 20.6μs (362% faster)

def test_invalid_test_framework():
    # Should raise assertion error for invalid test framework
    with pytest.raises(AssertionError):
        create_trace_replay_test_code(
            trace_file="trace.sqlite",
            functions_data=[],
            test_framework="nose",
            max_run_count=1,
        )

def test_function_with_missing_optional_keys():
    # Should handle missing class_name and function_properties gracefully
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            # class_name omitted
            "file_path": "/foo.py",
            "benchmark_function_name": "bar",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 96.9μs -> 22.9μs (324% faster)

def test_function_with_none_class_name():
    # Should treat None class_name as ""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": None,
            "file_path": "/foo.py",
            "benchmark_function_name": "bar",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 94.2μs -> 21.4μs (340% faster)

def test_multiple_methods_in_same_class():
    # Test multiple methods (including __init__) in same class
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/foo.py",
            "benchmark_function_name": "__init__",
            "function_properties": DummyFunctionProperties(),
        },
        {
            "module_name": "foo",
            "function_name": "doit",
            "class_name": "MyClass",
            "file_path": "/foo.py",
            "benchmark_function_name": "doit",
            "function_properties": DummyFunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=2,
    ); code = codeflash_output # 112μs -> 35.6μs (215% faster)

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

def test_large_number_of_functions():
    # Test with a large number of functions (500)
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": "",
            "file_path": f"/mod{i}.py",
            "benchmark_function_name": f"func{i}",
            "function_properties": DummyFunctionProperties(),
        }
        for i in range(500)
    ]
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 3.47ms -> 3.21ms (8.18% faster)
    # Check that all imports and test functions are present
    for i in range(500):
        pass

def test_large_number_of_methods_and_classes():
    # Test with 100 classes, each with 5 methods
    functions_data = []
    for i in range(100):
        for j in range(5):
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"method{j}",
                "class_name": f"Class{i}",
                "file_path": f"/mod{i}.py",
                "benchmark_function_name": f"method{j}",
                "function_properties": DummyFunctionProperties(),
            })
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output
    # Check that all class imports and test function names are present
    for i in range(100):
        for j in range(5):
            pass

def test_large_scale_performance():
    # Test that the function does not take too long for 1000 functions
    import time
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "class_name": "",
            "file_path": f"/mod{i}.py",
            "benchmark_function_name": f"func{i}",
            "function_properties": DummyFunctionProperties(),
        }
        for i in range(1000)
    ]
    start = time.time()
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output # 6.81ms -> 6.41ms (6.26% faster)
    elapsed = time.time() - start

def test_large_scale_class_and_static_methods():
    # Test with 50 classes, each with 5 static and 5 class methods
    functions_data = []
    for i in range(50):
        for j in range(5):
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"cmethod{j}",
                "class_name": f"Class{i}",
                "file_path": f"/mod{i}.py",
                "benchmark_function_name": f"cmethod{j}",
                "function_properties": DummyFunctionProperties(is_classmethod=True),
            })
            functions_data.append({
                "module_name": f"mod{i}",
                "function_name": f"smethod{j}",
                "class_name": f"Class{i}",
                "file_path": f"/mod{i}.py",
                "benchmark_function_name": f"smethod{j}",
                "function_properties": DummyFunctionProperties(is_staticmethod=True),
            })
    codeflash_output = create_trace_replay_test_code(
        trace_file="trace.sqlite",
        functions_data=functions_data,
        test_framework="pytest",
        max_run_count=1,
    ); code = codeflash_output
    # Check that all static/class methods are present
    for i in range(50):
        for j in range(5):
            pass
# 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-pr363-2025-06-22T22.41.59 and push.