⚡️ Speed up function create_trace_replay_test_code by 50% in PR #294 (add-timing-info-to-generated-tests)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #294

If you approve this dependent PR, these changes will be merged into the original PR branch add-timing-info-to-generated-tests.

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

---

📄 50% (0.50x) speedup for create_trace_replay_test_code in codeflash/benchmarking/replay_test.py

⏱️ Runtime : 11.5 milliseconds → 7.65 milliseconds (best of 348 runs)

📝 Explanation and details

Here's an optimized and faster version of your program. The main performance inefficiencies in the original code are.

• Repeated attribute accesses with dict.get() inside loops: Pre-collecting values boosts efficiency.
• Frequent string concatenations: Use f-strings carefully and only when necessary.
• Unnecessary use of sorted on a set each run. Build this directly from the data.
• Repeated construction of similar strings: Precompute or simplify where possible.
• Using textwrap.indent in a loop: Combine with minimal copies.
• No need for textwrap.dedent if formatting is already explicit.

Below is the refactored code following these optimizations.

Summary of the changes:

• Single pass for collecting imports and function names.
• Directly build up all test code as a list, for O(1) append performance and O(1) final string join.
• Minimized repeated calls to attribute-getting, string formatting, and function calls inside large loops.
• Efficient, manual indentation instead of textwrap.indent.
• Templates are constants, dedented only once.
• All constants precomputed outside the loop.

This will make your test code generation much faster and with much less memory overhead for large functions_data. No function signature or comments have been changed except for the relevant section reflecting the new optimized approach.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from __future__ import annotations

import textwrap
from typing import Any

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


# Mocks for function_properties
class DummyFunctionProps:
    def __init__(self, is_classmethod=False, is_staticmethod=False):
        self.is_classmethod = is_classmethod
        self.is_staticmethod = is_staticmethod

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

def test_single_function_pytest():
    """Test basic single function, pytest output."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.baz",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="pytest", max_run_count=10); result = codeflash_output # 90.7μs -> 13.3μs

def test_single_function_unittest():
    """Test basic single function, unittest output."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.baz",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="unittest", max_run_count=5); result = codeflash_output # 89.4μs -> 12.4μs

def test_multiple_functions_sorted():
    """Test multiple functions, sorted order in functions list."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.baz",
            "function_properties": DummyFunctionProps(),
        },
        {
            "module_name": "foo.bar",
            "function_name": "aaa",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.aaa",
            "function_properties": DummyFunctionProps(),
        },
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 94.9μs -> 16.9μs

def test_class_method():
    """Test a class method."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "doit",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.MyClass.doit",
            "function_properties": DummyFunctionProps(is_classmethod=True),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 89.6μs -> 14.6μs

def test_static_method():
    """Test a static method."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "stat",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.MyClass.stat",
            "function_properties": DummyFunctionProps(is_staticmethod=True),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 89.4μs -> 13.4μs

def test_regular_method():
    """Test a regular method (not static/class)."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "doit",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.MyClass.doit",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 92.0μs -> 15.3μs

def test_init_method():
    """Test __init__ method (should use *args[1:], **kwargs)."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.MyClass.__init__",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 90.0μs -> 14.0μs

def test_no_functions():
    """Test with empty functions_data."""
    functions_data = []
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 78.1μs -> 3.28μs

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

def test_module_with_multiple_dots():
    """Test module with multiple dots in name."""
    functions_data = [
        {
            "module_name": "foo.bar.baz.qux",
            "function_name": "zap",
            "file_path": "/tmp/foo/bar/baz/qux.py",
            "benchmark_function_name": "foo.bar.baz.qux.zap",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 88.3μs -> 12.8μs

def test_function_name_collision():
    """Test two functions with the same name but different modules."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "zap",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.zap",
            "function_properties": DummyFunctionProps(),
        },
        {
            "module_name": "foo.baz",
            "function_name": "zap",
            "file_path": "/tmp/foo/baz.py",
            "benchmark_function_name": "foo.baz.zap",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 93.9μs -> 16.6μs

def test_function_with_missing_optional_fields():
    """Test function dict missing optional fields (class_name)."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "zap",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.zap",
            "function_properties": DummyFunctionProps(),
            # no class_name
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 87.5μs -> 11.3μs

def test_function_with_weird_names():
    """Test function and class names with underscores and digits."""
    functions_data = [
        {
            "module_name": "foo.bar2",
            "function_name": "zap_123",
            "class_name": "Class_456",
            "file_path": "/tmp/foo/bar2.py",
            "benchmark_function_name": "foo.bar2.Class_456.zap_123",
            "function_properties": DummyFunctionProps(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 92.5μs -> 15.7μs

def test_assert_invalid_framework():
    """Test that invalid test_framework raises AssertionError."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.baz",
            "function_properties": DummyFunctionProps(),
        }
    ]
    with pytest.raises(AssertionError):
        create_trace_replay_test_code("/tmp/trace.sqlite", functions_data, test_framework="nose")

def test_function_with_init_and_others():
    """Test that __init__ is not included in 'functions' list, but test is generated."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "__init__",
            "class_name": "MyClass",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.MyClass.__init__",
            "function_properties": DummyFunctionProps(),
        },
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "foo.bar.baz",
            "function_properties": DummyFunctionProps(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("/tmp/trace.sqlite", functions_data); result = codeflash_output # 99.5μs -> 22.0μs

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

def test_large_number_of_functions():
    """Test with a large number of functions (scalability)."""
    functions_data = [
        {
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"mod{i}.func{i}",
            "function_properties": DummyFunctionProps(),
        }
        for i in range(100)
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 477μs -> 334μs
    # Should contain all test functions
    for i in range(100):
        pass

def test_large_number_of_methods_and_classes():
    """Test with many classes and methods."""
    functions_data = [
        {
            "module_name": f"mod{i//10}",
            "function_name": f"meth{i%10}",
            "class_name": f"Cls{i//10}",
            "file_path": f"/tmp/mod{i//10}.py",
            "benchmark_function_name": f"mod{i//10}.Cls{i//10}.meth{i%10}",
            "function_properties": DummyFunctionProps(),
        }
        for i in range(100)
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 728μs -> 546μs
    # Should contain all class imports and test functions
    for i in range(10):
        pass
    for i in range(100):
        pass

def test_large_scale_with_mixed_types():
    """Test with a mix of functions, static, class, and regular methods."""
    functions_data = []
    for i in range(30):
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"func{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"mod{i}.func{i}",
            "function_properties": DummyFunctionProps(),
        })
    for i in range(30, 60):
        functions_data.append({
            "module_name": f"mod{i//2}",
            "function_name": f"meth{i}",
            "class_name": f"Cls{i//2}",
            "file_path": f"/tmp/mod{i//2}.py",
            "benchmark_function_name": f"mod{i//2}.Cls{i//2}.meth{i}",
            "function_properties": DummyFunctionProps(is_classmethod=True),
        })
    for i in range(60, 90):
        functions_data.append({
            "module_name": f"mod{i//3}",
            "function_name": f"meth{i}",
            "class_name": f"Cls{i//3}",
            "file_path": f"/tmp/mod{i//3}.py",
            "benchmark_function_name": f"mod{i//3}.Cls{i//3}.meth{i}",
            "function_properties": DummyFunctionProps(is_staticmethod=True),
        })
    for i in range(90, 100):
        functions_data.append({
            "module_name": f"mod{i}",
            "function_name": f"meth{i}",
            "class_name": f"Cls{i}",
            "file_path": f"/tmp/mod{i}.py",
            "benchmark_function_name": f"mod{i}.Cls{i}.meth{i}",
            "function_properties": DummyFunctionProps(),
        })
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data); result = codeflash_output # 579μs -> 427μs
# 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 re
import sys
import textwrap
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 to simulate function_properties attribute in functions_data
class FunctionProperties:
    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_basic_single_function_pytest():
    """Test with a single standalone function, pytest framework."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz",
            "function_properties": FunctionProperties(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="pytest"); code = codeflash_output # 89.6μs -> 13.4μs

def test_basic_single_function_unittest():
    """Test with a single standalone function, unittest framework."""
    functions_data = [
        {
            "module_name": "foo.bar",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo/bar.py",
            "benchmark_function_name": "baz",
            "function_properties": FunctionProperties(),
        }
    ]
    trace_file = "/tmp/trace.sqlite"
    codeflash_output = create_trace_replay_test_code(trace_file, functions_data, test_framework="unittest"); code = codeflash_output # 88.0μs -> 12.5μs

def test_basic_multiple_functions():
    """Test with multiple standalone functions, ensure all are present."""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "f1",
            "class_name": None,
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "f1",
            "function_properties": FunctionProperties(),
        },
        {
            "module_name": "foo",
            "function_name": "f2",
            "class_name": None,
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "f2",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 94.0μs -> 16.3μs

def test_basic_method_and_function():
    """Test with a standalone function and a method (not class/static), pytest framework."""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "f",
            "class_name": None,
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "f",
            "function_properties": FunctionProperties(),
        },
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": "Baz",
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "bar",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 98.6μs -> 20.3μs

def test_basic_classmethod_and_staticmethod():
    """Test with a classmethod and a staticmethod."""
    functions_data = [
        {
            "module_name": "mod",
            "function_name": "cm",
            "class_name": "C",
            "file_path": "/tmp/mod.py",
            "benchmark_function_name": "cm",
            "function_properties": FunctionProperties(is_classmethod=True),
        },
        {
            "module_name": "mod",
            "function_name": "sm",
            "class_name": "C",
            "file_path": "/tmp/mod.py",
            "benchmark_function_name": "sm",
            "function_properties": FunctionProperties(is_staticmethod=True),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 97.3μs -> 20.1μs

def test_basic_init_method():
    """Test with a __init__ method, ensure correct handling."""
    functions_data = [
        {
            "module_name": "pkg.mod",
            "function_name": "__init__",
            "class_name": "Foo",
            "file_path": "/tmp/pkg/mod.py",
            "benchmark_function_name": "__init__",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 90.2μs -> 14.3μs

# 2. EDGE TEST CASES

def test_empty_functions_data():
    """Test with empty functions_data, should still return code with imports and empty functions list."""
    codeflash_output = create_trace_replay_test_code("trace.sqlite", [], test_framework="pytest"); code = codeflash_output # 77.5μs -> 3.62μs

def test_function_name_init_is_skipped_in_functions_list():
    """Test that __init__ is excluded from the functions list."""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "__init__",
            "class_name": "Bar",
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "__init__",
            "function_properties": FunctionProperties(),
        },
        {
            "module_name": "foo",
            "function_name": "baz",
            "class_name": None,
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "baz",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 99.9μs -> 20.9μs

def test_function_properties_missing_keys():
    """Test that missing class_name or function_properties does not raise, but uses defaults."""
    # class_name missing
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "baz",
            # "class_name" omitted
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "baz",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 88.3μs -> 11.7μs
    # Should not raise

def test_non_ascii_module_and_function_names():
    """Test with non-ASCII module, class, and function names."""
    functions_data = [
        {
            "module_name": "módülé",
            "function_name": "fünç",
            "class_name": "Çláss",
            "file_path": "/tmp/üñîçødë.py",
            "benchmark_function_name": "fünç",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("üñîçødë.sqlite", functions_data); code = codeflash_output # 95.0μs -> 18.1μs

def test_max_run_count_parameter():
    """Test that max_run_count is respected in the generated test code."""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "bar",
            "class_name": None,
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "bar",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data, max_run_count=17); code = codeflash_output # 88.8μs -> 11.9μs

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

def test_class_and_function_same_name():
    """Test when class and function have the same name, aliasing is correct."""
    functions_data = [
        {
            "module_name": "foo",
            "function_name": "Bar",
            "class_name": "Bar",
            "file_path": "/tmp/foo.py",
            "benchmark_function_name": "Bar",
            "function_properties": FunctionProperties(),
        }
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 92.6μs -> 15.7μs

# 3. LARGE SCALE TEST CASES

def test_large_number_of_functions():
    """Test with a large number (500) of standalone functions."""
    N = 500
    functions_data = [
        {
            "module_name": "bigmod",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": "/tmp/bigmod.py",
            "benchmark_function_name": f"func{i}",
            "function_properties": FunctionProperties(),
        }
        for i in range(N)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 477μs -> 334μs
    # All imports and test functions should be present
    for i in [0, N//2, N-1]:
        pass
    # Functions list should have all function names sorted
    expected_list = [f"func{i}" for i in range(N)]

def test_large_number_of_methods():
    """Test with 300 methods belonging to different classes."""
    N = 300
    functions_data = [
        {
            "module_name": "mod",
            "function_name": f"m{i}",
            "class_name": f"C{i%10}",
            "file_path": "/tmp/mod.py",
            "benchmark_function_name": f"m{i}",
            "function_properties": FunctionProperties(),
        }
        for i in range(N)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 2.05ms -> 1.64ms
    # Check a few representative imports and test function names
    for i in [0, 50, 299]:
        class_name = f"C{i%10}"

def test_large_mixed_functions_and_methods():
    """Test with a mix of 100 functions and 100 methods."""
    N = 100
    functions_data = [
        {
            "module_name": "mix",
            "function_name": f"f{i}",
            "class_name": None,
            "file_path": "/tmp/mix.py",
            "benchmark_function_name": f"f{i}",
            "function_properties": FunctionProperties(),
        }
        for i in range(N)
    ] + [
        {
            "module_name": "mix",
            "function_name": f"m{i}",
            "class_name": "C",
            "file_path": "/tmp/mix.py",
            "benchmark_function_name": f"m{i}",
            "function_properties": FunctionProperties(),
        }
        for i in range(N)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 1.14ms -> 887μs
    # Check some function and method presence
    for i in [0, 50, 99]:
        pass

def test_large_scale_performance():
    """Test that large input does not exceed reasonable code size and is efficient."""
    N = 500
    functions_data = [
        {
            "module_name": "perf",
            "function_name": f"func{i}",
            "class_name": None,
            "file_path": "/tmp/perf.py",
            "benchmark_function_name": f"func{i}",
            "function_properties": FunctionProperties(),
        }
        for i in range(N)
    ]
    codeflash_output = create_trace_replay_test_code("trace.sqlite", functions_data); code = codeflash_output # 2.05ms -> 1.64ms
# 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-pr294-2025-06-06T03.43.25 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #294 by misrasaurabh1 was closed.