⚡️ Speed up function remove_colorstr by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for remove_colorstr in ultralytics/utils/__init__.py

⏱️ Runtime : 516 microseconds → 465 microseconds (best of 250 runs)

📝 Explanation and details

The optimization moves the regex compilation from inside the function to module level, eliminating redundant regex compilation on every function call.

Key Changes:

• Pre-compiled regex: _ANSI_ESCAPE_RE is compiled once at module import time instead of being recompiled in every function call
• Type hints added: Function signature now includes explicit type annotations for better code clarity

Why This Optimization Works:
The original code compiled the same regex pattern r"\x1B\[[0-9;]*[A-Za-z]" on every function invocation, which is expensive. The line profiler shows that 40.5% of the original function's time was spent on re.compile(). By moving this compilation to module level, we eliminate this overhead entirely.

Performance Impact:

• 11% overall speedup with consistent improvements across all test cases
• Best performance gains: Simple cases like empty strings and plain text see 100%+ speedup (1.76μs → 777ns)
• Moderate gains for complex cases: Unicode and large-scale tests show smaller but consistent 2-35% improvements
• Large-scale workloads: Even with thousands of ANSI codes, the optimization maintains 1-8% performance gains

Real-World Benefits:
Based on the function reference, remove_colorstr() is called from hyperparameter tuning code where it processes log headers that likely contain colored terminal output. Since tuning involves many iterations with logging, this 11% speedup in string processing can accumulate to meaningful time savings during long optimization runs. The function appears to be in a hot path where logging occurs frequently, making this micro-optimization worthwhile.

The optimization is particularly effective for workloads with repeated calls to remove ANSI codes from strings, which is common in CLI tools and logging systems where colored output needs to be cleaned for file storage or plain text display.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# imports
from ultralytics.utils.__init__ import remove_colorstr

# unit tests

# ------------------------
# BASIC TEST CASES
# ------------------------


def test_remove_colorstr_with_plain_text():
    # No ANSI codes; should return unchanged string
    codeflash_output = remove_colorstr("hello world")  # 1.76μs -> 777ns (127% faster)


def test_remove_colorstr_with_single_ansi_code():
    # One ANSI code at the start
    colored = "\x1b[31mhello world"
    codeflash_output = remove_colorstr(colored)  # 2.06μs -> 1.13μs (82.5% faster)


def test_remove_colorstr_with_multiple_ansi_codes():
    # Multiple codes at different positions
    colored = "\x1b[31mhello \x1b[1mworld\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 2.30μs -> 1.46μs (57.9% faster)


def test_remove_colorstr_with_reset_code():
    # Reset code at the end
    colored = "hello world\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 1.78μs -> 1.01μs (77.1% faster)


def test_remove_colorstr_with_embedded_code():
    # ANSI code in the middle of text
    colored = "hello \x1b[34mworld"
    codeflash_output = remove_colorstr(colored)  # 2.00μs -> 1.12μs (77.5% faster)


def test_remove_colorstr_with_multiple_codes_and_text():
    # Multiple codes interspersed
    colored = "\x1b[31mred\x1b[0m and \x1b[32mgreen\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 2.48μs -> 1.73μs (43.9% faster)


# ------------------------
# EDGE TEST CASES
# ------------------------


def test_remove_colorstr_empty_string():
    # Edge: empty string input
    codeflash_output = remove_colorstr("")  # 1.42μs -> 632ns (125% faster)


def test_remove_colorstr_only_ansi_codes():
    # Edge: input is only ANSI codes, no text
    colored = "\x1b[31m\x1b[0m\x1b[1m"
    codeflash_output = remove_colorstr(colored)  # 2.10μs -> 1.32μs (59.5% faster)


def test_remove_colorstr_ansi_code_at_end():
    # Edge: ANSI code at the very end
    colored = "text\x1b[31m"
    codeflash_output = remove_colorstr(colored)  # 1.79μs -> 1.05μs (70.9% faster)


def test_remove_colorstr_ansi_code_at_start_and_end():
    # Edge: codes at both start and end
    colored = "\x1b[31mtext\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 1.99μs -> 1.18μs (69.2% faster)


def test_remove_colorstr_adjacent_ansi_codes():
    # Edge: multiple adjacent codes
    colored = "\x1b[31m\x1b[1mhello"
    codeflash_output = remove_colorstr(colored)  # 2.08μs -> 1.26μs (65.5% faster)


def test_remove_colorstr_nested_codes():
    # Edge: codes nested within text
    colored = "A\x1b[31mB\x1b[0mC"
    codeflash_output = remove_colorstr(colored)  # 2.08μs -> 1.29μs (61.3% faster)


def test_remove_colorstr_with_non_ansi_escape_sequences():
    # Edge: string with similar but non-ANSI sequences should not be removed
    colored = "hello \x1b[XYZm world"
    # The regex only matches valid ANSI codes; this should remain unchanged
    codeflash_output = remove_colorstr(colored)  # 1.88μs -> 1.17μs (60.5% faster)


def test_remove_colorstr_with_unicode_characters():
    # Edge: string with unicode and ANSI codes
    colored = "\x1b[31m你好，世界\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 3.14μs -> 2.33μs (34.5% faster)


def test_remove_colorstr_with_incomplete_ansi_code():
    # Edge: incomplete ANSI code should not be removed
    colored = "hello\x1b[31"
    codeflash_output = remove_colorstr(colored)  # 1.93μs -> 1.08μs (78.8% faster)


def test_remove_colorstr_with_semicolon_parameters():
    # Edge: ANSI code with multiple parameters
    colored = "\x1b[1;31;40mhello"
    codeflash_output = remove_colorstr(colored)  # 2.04μs -> 1.21μs (68.5% faster)


def test_remove_colorstr_with_uppercase_final_letter():
    # Edge: ANSI codes with uppercase final letter
    colored = "\x1b[31Mhello"
    codeflash_output = remove_colorstr(colored)  # 1.91μs -> 1.08μs (77.1% faster)


def test_remove_colorstr_with_mixed_case_final_letter():
    # Edge: ANSI codes with lowercase and uppercase final letter
    colored = "\x1b[31mHello\x1b[32MWorld"
    codeflash_output = remove_colorstr(colored)  # 2.10μs -> 1.33μs (57.8% faster)


# ------------------------
# LARGE SCALE TEST CASES
# ------------------------


def test_remove_colorstr_large_text_with_many_codes():
    # Large: text of 1000 words, each word colored
    words = ["word{}".format(i) for i in range(1000)]
    colored = "".join("\x1b[{}m{}".format(30 + (i % 8), w) for i, w in enumerate(words))
    expected = "".join(words)
    codeflash_output = remove_colorstr(colored)  # 75.4μs -> 73.8μs (2.16% faster)


def test_remove_colorstr_large_text_with_interspersed_codes():
    # Large: 1000 words, every other word colored
    words = ["word{}".format(i) for i in range(1000)]
    colored = ""
    for i, w in enumerate(words):
        if i % 2 == 0:
            colored += "\x1b[31m" + w
        else:
            colored += w
    expected = "".join(words)
    codeflash_output = remove_colorstr(colored)  # 37.7μs -> 36.7μs (2.72% faster)


def test_remove_colorstr_large_text_no_codes():
    # Large: 1000 words, no ANSI codes
    words = ["word{}".format(i) for i in range(1000)]
    text = "".join(words)
    codeflash_output = remove_colorstr(text)  # 2.58μs -> 1.74μs (48.3% faster)


def test_remove_colorstr_large_text_all_codes():
    # Large: 1000 consecutive ANSI codes, no text
    colored = "".join("\x1b[{}m".format(30 + (i % 8)) for i in range(1000))
    codeflash_output = remove_colorstr(colored)  # 55.3μs -> 53.3μs (3.86% faster)


def test_remove_colorstr_large_text_mixed_content():
    # Large: 1000 elements, alternating text and ANSI codes
    colored = ""
    expected = ""
    for i in range(1000):
        if i % 2 == 0:
            colored += "\x1b[31m"
        else:
            colored += "word{}".format(i)
            expected += "word{}".format(i)
    codeflash_output = remove_colorstr(colored)  # 36.4μs -> 35.4μs (2.86% faster)


# ------------------------
# ADDITIONAL EDGE CASES
# ------------------------


def test_remove_colorstr_with_newlines_and_tabs():
    # Edge: string with newlines and tabs and ANSI codes
    colored = "\x1b[31mhello\nworld\t\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 2.19μs -> 1.32μs (66.0% faster)


def test_remove_colorstr_with_multiple_codes_in_a_row():
    # Edge: several codes in a row, no text between
    colored = "\x1b[31m\x1b[1m\x1b[4mtext"
    codeflash_output = remove_colorstr(colored)  # 2.22μs -> 1.43μs (54.8% faster)


def test_remove_colorstr_with_long_code_parameters():
    # Edge: ANSI code with long parameter list
    colored = "\x1b[1;2;3;4;5;6;7;8;9;10mhello"
    codeflash_output = remove_colorstr(colored)  # 2.07μs -> 1.25μs (65.6% faster)


def test_remove_colorstr_with_mixed_valid_and_invalid_codes():
    # Edge: valid and invalid codes mixed
    colored = "\x1b[31mhello\x1b[XYZmworld"
    # Only the valid code is removed
    codeflash_output = remove_colorstr(colored)  # 2.16μs -> 1.39μs (56.0% faster)


def test_remove_colorstr_with_code_followed_by_non_letter():
    # Edge: code not ending with a letter should not be removed
    colored = "hello\x1b[31"
    codeflash_output = remove_colorstr(colored)  # 1.90μs -> 1.08μs (75.8% faster)


def test_remove_colorstr_with_code_embedded_in_word():
    # Edge: code embedded within a word
    colored = "he\x1b[31mllo"
    codeflash_output = remove_colorstr(colored)  # 2.04μs -> 1.17μs (74.5% faster)


def test_remove_colorstr_with_multiple_types_of_codes():
    # Edge: codes ending with different letters
    colored = "\x1b[31mred\x1b[32Kgreen\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 2.27μs -> 1.50μs (51.8% faster)


def test_remove_colorstr_with_spaces_around_codes():
    # Edge: codes surrounded by spaces
    colored = " \x1b[31m hello \x1b[0m "
    codeflash_output = remove_colorstr(colored)  # 2.13μs -> 1.37μs (55.2% faster)


def test_remove_colorstr_with_realistic_colored_log():
    # Edge: realistic log line with color codes
    colored = "[INFO] \x1b[32mSuccess\x1b[0m: \x1b[31mError\x1b[0m"
    codeflash_output = remove_colorstr(colored)  # 2.42μs -> 1.51μs (60.1% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import random
import string

# imports
import pytest  # used for our unit tests
from ultralytics.utils.__init__ import remove_colorstr

# unit tests

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


def test_remove_colorstr_basic_no_ansi():
    # No ANSI codes, should return unchanged
    s = "hello world"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.47μs -> 720ns (104% faster)


def test_remove_colorstr_basic_single_ansi():
    # Single ANSI code at start
    s = "\x1b[31mhello world"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.98μs -> 1.12μs (76.0% faster)


def test_remove_colorstr_basic_ansi_in_middle():
    # ANSI code in the middle
    s = "hello \x1b[1mworld"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.87μs -> 1.13μs (65.4% faster)


def test_remove_colorstr_basic_multiple_ansi():
    # Multiple ANSI codes
    s = "\x1b[1mhello\x1b[0m world\x1b[32m!"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 2.29μs -> 1.54μs (48.7% faster)


def test_remove_colorstr_basic_ansi_with_semicolons():
    # ANSI code with multiple parameters
    s = "\x1b[1;31mhello world"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.97μs -> 1.09μs (81.0% faster)


def test_remove_colorstr_basic_ansi_at_end():
    # ANSI code at the end
    s = "hello world\x1b[0m"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.80μs -> 1.05μs (71.6% faster)


def test_remove_colorstr_basic_ansi_only():
    # String is only ANSI code
    s = "\x1b[0m"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.75μs -> 956ns (83.2% faster)


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


def test_remove_colorstr_empty_string():
    # Empty string input
    s = ""
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.41μs -> 621ns (127% faster)


def test_remove_colorstr_ansi_adjacent():
    # Adjacent ANSI codes
    s = "\x1b[1m\x1b[32mhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 2.14μs -> 1.32μs (61.9% faster)


def test_remove_colorstr_ansi_with_unusual_letters():
    # ANSI code ending in unusual letter
    s = "\x1b[31Khello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.94μs -> 1.04μs (85.3% faster)


def test_remove_colorstr_ansi_with_no_parameters():
    # ANSI code with no parameters
    s = "\x1b[mhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.86μs -> 978ns (90.0% faster)


def test_remove_colorstr_ansi_in_between_chars():
    # ANSI code between characters
    s = "he\x1b[1mllo"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.97μs -> 1.14μs (73.2% faster)


def test_remove_colorstr_ansi_with_extra_escape():
    # Extra ESC character not followed by [
    s = "\x1bhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.49μs -> 670ns (123% faster)


def test_remove_colorstr_non_string_input_raises():
    # Non-string input should raise TypeError
    with pytest.raises(TypeError):
        remove_colorstr(123)  # 2.33μs -> 1.46μs (59.6% faster)


def test_remove_colorstr_ansi_with_long_params():
    # ANSI code with long parameter list
    s = "\x1b[1;2;3;4;5;6;7;8;9;10mhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 2.18μs -> 1.32μs (65.5% faster)


def test_remove_colorstr_malformed_ansi_code():
    # Malformed ANSI code (missing [)
    s = "\x1b31mhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.46μs -> 668ns (119% faster)


def test_remove_colorstr_ansi_with_multiple_letters():
    # ANSI code ending with two letters (invalid, should only match one letter)
    s = "\x1b[31mmhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.91μs -> 1.17μs (63.5% faster)


def test_remove_colorstr_unicode_string():
    # Unicode characters mixed with ANSI
    s = "\x1b[31m你好，世界\x1b[0m"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 3.18μs -> 2.49μs (27.8% faster)


def test_remove_colorstr_surrounding_ansi():
    # ANSI codes surrounding text
    s = "\x1b[32mhello\x1b[0m"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 2.06μs -> 1.33μs (54.4% faster)


def test_remove_colorstr_ansi_with_nonstandard_letter():
    # ANSI code ending with nonstandard letter (e.g., 'x')
    s = "\x1b[31xhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.88μs -> 1.09μs (72.2% faster)


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


def test_remove_colorstr_large_text_no_ansi():
    # Large text, no ANSI codes
    s = "a" * 1000
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 1.59μs -> 791ns (101% faster)


def test_remove_colorstr_large_text_with_ansi_everywhere():
    # Large text with ANSI codes inserted every 10 characters
    base = "abcdefghij"
    s = ""
    for i in range(100):
        s += "\x1b[1;31m" + base
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 11.2μs -> 10.5μs (6.27% faster)


def test_remove_colorstr_large_text_with_random_ansi():
    # Large random text with random ANSI codes
    chars = string.ascii_letters + string.digits
    s = []
    for i in range(500):
        # Randomly insert ANSI code
        if random.random() < 0.2:
            s.append("\x1b[1;3{}m".format(random.randint(0, 9)))
        s.append(random.choice(chars))
    s = "".join(s)
    # Remove all ANSI codes
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 35.6μs -> 34.0μs (4.68% faster)


def test_remove_colorstr_large_text_all_ansi():
    # All characters are ANSI codes
    s = "".join(["\x1b[{}m".format(i) for i in range(1000)])
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 54.7μs -> 54.0μs (1.40% faster)


def test_remove_colorstr_large_text_mixed():
    # Large text, mixed with and without ANSI codes
    s = ""
    for i in range(500):
        if i % 2 == 0:
            s += "\x1b[1mhello"
        else:
            s += "world"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 19.7μs -> 18.7μs (5.49% faster)


def test_remove_colorstr_large_text_unicode_and_ansi():
    # Large text with unicode and ANSI codes
    s = ""
    for i in range(500):
        s += "\x1b[31m你好世界"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 46.4μs -> 46.0μs (0.747% faster)


def test_remove_colorstr_large_text_edge_ansi():
    # Large text with edge-case ANSI codes (no parameters)
    s = ""
    for i in range(500):
        s += "\x1b[mhello"
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 33.4μs -> 32.7μs (1.91% faster)


def test_remove_colorstr_performance_large_string():
    # Performance test: large string with many ANSI codes
    s = "hello\x1b[31mworld" * 100
    codeflash_output = remove_colorstr(s)
    result = codeflash_output  # 10.6μs -> 9.79μs (7.90% 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-remove_colorstr-mi8cu2ou and push.