⚡️ Speed up function remove_trailing_yaml by 12%

[codeflash-ai]

📄 12% (0.12x) speedup for remove_trailing_yaml in higgsfield/internal/experiment/params.py

⏱️ Runtime : 29.2 microseconds → 26.0 microseconds (best of 489 runs)

📝 Explanation and details

The optimization eliminates unnecessary variable assignment and function call overhead by directly inlining the string literal and its length.

Key changes:

1. Removed variable assignment: Instead of storing "\n...\n" in trailing_yaml, the string literal is used directly in the endswith() call
2. Hardcoded slice index: Replaced s[: -len(trailing_yaml)] with s[:-5], eliminating the len() function call

Why this is faster:

• Eliminates the overhead of variable creation and assignment (29.4% of original runtime)
• Removes a function call to len() which was being executed on every match
• The Python interpreter can optimize direct string literals better than variables

Performance characteristics:
The optimization shows consistent speedups across all test scenarios, with the most significant gains (20-30%) occurring when the trailing YAML pattern is actually found and removed. Even when the pattern isn't present, there's still a modest improvement (1-10%) from eliminating the variable assignment. The optimization is particularly effective for larger inputs with trailing patterns, where the len() call elimination provides more substantial benefits.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from higgsfield.internal.experiment.params import remove_trailing_yaml

# unit tests

# 1. Basic Test Cases

def test_remove_trailing_yaml_basic_removal():
    # Should remove trailing '\n...\n' from the end
    input_str = "key: value\n...\n"
    expected = "key: value"
    codeflash_output = remove_trailing_yaml(input_str) # 823ns -> 647ns (27.2% faster)

def test_remove_trailing_yaml_no_trailing_yaml():
    # Should return unchanged string if no trailing '\n...\n'
    input_str = "key: value\nfoo: bar"
    expected = "key: value\nfoo: bar"
    codeflash_output = remove_trailing_yaml(input_str) # 506ns -> 459ns (10.2% faster)

def test_remove_trailing_yaml_multiple_yaml_not_at_end():
    # Only removes if '\n...\n' is at the end
    input_str = "key: value\n...\nfoo: bar\n"
    expected = "key: value\n...\nfoo: bar\n"
    codeflash_output = remove_trailing_yaml(input_str) # 464ns -> 453ns (2.43% faster)

def test_remove_trailing_yaml_empty_string():
    # Should return empty string unchanged
    input_str = ""
    expected = ""
    codeflash_output = remove_trailing_yaml(input_str) # 446ns -> 480ns (7.08% slower)

def test_remove_trailing_yaml_exact_trailing_yaml():
    # Input is exactly the trailing yaml
    input_str = "\n...\n"
    expected = ""
    codeflash_output = remove_trailing_yaml(input_str) # 877ns -> 708ns (23.9% faster)

def test_remove_trailing_yaml_trailing_yaml_with_extra_newline():
    # Should not remove if extra newline after trailing yaml
    input_str = "key: value\n...\n\n"
    expected = "key: value\n...\n\n"
    codeflash_output = remove_trailing_yaml(input_str) # 503ns -> 454ns (10.8% faster)

def test_remove_trailing_yaml_trailing_yaml_with_spaces():
    # Should not remove if trailing yaml has spaces
    input_str = "key: value\n... \n"
    expected = "key: value\n... \n"
    codeflash_output = remove_trailing_yaml(input_str) # 476ns -> 479ns (0.626% slower)

# 2. Edge Test Cases

def test_remove_trailing_yaml_trailing_yaml_in_middle():
    # Should not remove if '\n...\n' is not at the end
    input_str = "key: value\n...\nfoo: bar\n"
    expected = "key: value\n...\nfoo: bar\n"
    codeflash_output = remove_trailing_yaml(input_str) # 476ns -> 436ns (9.17% faster)

def test_remove_trailing_yaml_multiple_trailing_yaml_at_end():
    # Only removes one trailing yaml
    input_str = "key: value\n...\n...\n"
    expected = "key: value\n...\n"
    codeflash_output = remove_trailing_yaml(input_str) # 897ns -> 714ns (25.6% faster)

def test_remove_trailing_yaml_trailing_yaml_with_tabs():
    # Should not remove if trailing yaml has tabs
    input_str = "key: value\n...\n\t"
    expected = "key: value\n...\n\t"
    codeflash_output = remove_trailing_yaml(input_str) # 479ns -> 474ns (1.05% faster)

def test_remove_trailing_yaml_unicode_characters():
    # Should correctly handle unicode characters
    input_str = "ключ: значение\n...\n"
    expected = "ключ: значение"
    codeflash_output = remove_trailing_yaml(input_str) # 1.21μs -> 1.03μs (17.6% faster)

def test_remove_trailing_yaml_only_newlines():
    # String with only newlines should not be changed
    input_str = "\n\n"
    expected = "\n\n"
    codeflash_output = remove_trailing_yaml(input_str) # 472ns -> 463ns (1.94% faster)

def test_remove_trailing_yaml_partial_yaml_at_end():
    # Should not remove if only part of the trailing yaml is present
    input_str = "key: value\n...\nfoo"
    expected = "key: value\n...\nfoo"
    codeflash_output = remove_trailing_yaml(input_str) # 500ns -> 449ns (11.4% faster)

def test_remove_trailing_yaml_trailing_yaml_with_crlf():
    # Should not remove if line endings are CRLF
    input_str = "key: value\r\n...\r\n"
    expected = "key: value\r\n...\r\n"
    codeflash_output = remove_trailing_yaml(input_str) # 450ns -> 456ns (1.32% slower)

def test_remove_trailing_yaml_trailing_yaml_with_extra_content():
    # Should not remove if there is extra content after trailing yaml
    input_str = "key: value\n...\nextra"
    expected = "key: value\n...\nextra"
    codeflash_output = remove_trailing_yaml(input_str) # 452ns -> 456ns (0.877% slower)

# 3. Large Scale Test Cases

def test_remove_trailing_yaml_large_text_with_trailing_yaml():
    # Large input with trailing yaml at the end
    base = "key: value\n" * 500
    input_str = base + "\n...\n"
    expected = base
    codeflash_output = remove_trailing_yaml(input_str) # 1.12μs -> 931ns (20.4% faster)

def test_remove_trailing_yaml_large_text_without_trailing_yaml():
    # Large input without trailing yaml
    base = "key: value\n" * 999
    input_str = base
    expected = base
    codeflash_output = remove_trailing_yaml(input_str) # 483ns -> 440ns (9.77% faster)

def test_remove_trailing_yaml_large_text_with_multiple_trailing_yaml():
    # Large input with multiple trailing yaml at the end
    base = "key: value\n" * 998
    input_str = base + "\n...\n...\n"
    expected = base + "\n...\n"
    codeflash_output = remove_trailing_yaml(input_str) # 1.39μs -> 1.25μs (11.5% faster)

def test_remove_trailing_yaml_large_text_with_trailing_yaml_and_extra_newline():
    # Large input with trailing yaml but extra newline after
    base = "key: value\n" * 999
    input_str = base + "\n...\n\n"
    expected = base + "\n...\n\n"
    codeflash_output = remove_trailing_yaml(input_str) # 477ns -> 457ns (4.38% faster)

def test_remove_trailing_yaml_large_text_with_trailing_yaml_and_extra_content():
    # Large input with trailing yaml and extra content after
    base = "key: value\n" * 999
    input_str = base + "\n...\nextra"
    expected = base + "\n...\nextra"
    codeflash_output = remove_trailing_yaml(input_str) # 481ns -> 428ns (12.4% 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 higgsfield.internal.experiment.params import remove_trailing_yaml

# unit tests

# -------------------------
# 1. Basic Test Cases
# -------------------------

def test_remove_trailing_yaml_basic_present():
    # Basic: String ends with trailing YAML marker
    input_str = "foo\nbar\n...\n"
    expected = "foo\nbar"
    codeflash_output = remove_trailing_yaml(input_str) # 903ns -> 718ns (25.8% faster)

def test_remove_trailing_yaml_basic_absent():
    # Basic: String does not end with marker
    input_str = "foo\nbar\nbaz"
    expected = "foo\nbar\nbaz"
    codeflash_output = remove_trailing_yaml(input_str) # 514ns -> 504ns (1.98% faster)

def test_remove_trailing_yaml_basic_marker_in_middle():
    # Basic: Marker present in middle, not at end
    input_str = "foo\n...\nbar\n"
    expected = "foo\n...\nbar\n"
    codeflash_output = remove_trailing_yaml(input_str) # 512ns -> 483ns (6.00% faster)

def test_remove_trailing_yaml_basic_only_marker():
    # Basic: String is exactly the marker
    input_str = "\n...\n"
    expected = ""
    codeflash_output = remove_trailing_yaml(input_str) # 881ns -> 769ns (14.6% faster)

def test_remove_trailing_yaml_basic_marker_with_trailing_space():
    # Basic: Marker with trailing space (should not be removed)
    input_str = "foo\nbar\n...\n "
    expected = "foo\nbar\n...\n "
    codeflash_output = remove_trailing_yaml(input_str) # 527ns -> 489ns (7.77% faster)

# -------------------------
# 2. Edge Test Cases
# -------------------------

def test_remove_trailing_yaml_empty_string():
    # Edge: Empty string input
    input_str = ""
    expected = ""
    codeflash_output = remove_trailing_yaml(input_str) # 499ns -> 489ns (2.04% faster)

def test_remove_trailing_yaml_marker_at_start():
    # Edge: Marker at start of string
    input_str = "\n...\nfoo\nbar"
    expected = "\n...\nfoo\nbar"
    codeflash_output = remove_trailing_yaml(input_str) # 499ns -> 482ns (3.53% faster)

def test_remove_trailing_yaml_marker_multiple_times_at_end():
    # Edge: Multiple markers at end, only one should be removed
    input_str = "foo\nbar\n...\n...\n"
    expected = "foo\nbar\n..."
    codeflash_output = remove_trailing_yaml(input_str) # 889ns -> 685ns (29.8% faster)

def test_remove_trailing_yaml_marker_with_extra_newline():
    # Edge: Marker with extra newline at end (should not be removed)
    input_str = "foo\nbar\n...\n\n"
    expected = "foo\nbar\n...\n\n"
    codeflash_output = remove_trailing_yaml(input_str) # 497ns -> 476ns (4.41% faster)

def test_remove_trailing_yaml_marker_with_missing_newline():
    # Edge: Marker missing leading newline (should not be removed)
    input_str = "foo\nbar...\n"
    expected = "foo\nbar...\n"
    codeflash_output = remove_trailing_yaml(input_str) # 502ns -> 441ns (13.8% faster)

def test_remove_trailing_yaml_marker_with_extra_characters():
    # Edge: Marker followed by extra characters (should not be removed)
    input_str = "foo\nbar\n...\nabc"
    expected = "foo\nbar\n...\nabc"
    codeflash_output = remove_trailing_yaml(input_str) # 446ns -> 458ns (2.62% slower)

def test_remove_trailing_yaml_only_partial_marker_at_end():
    # Edge: Partial marker at end (should not be removed)
    input_str = "foo\nbar\n..\n"
    expected = "foo\nbar\n..\n"
    codeflash_output = remove_trailing_yaml(input_str) # 471ns -> 455ns (3.52% faster)

def test_remove_trailing_yaml_unicode_characters():
    # Edge: Unicode characters present
    input_str = "foo\u2603\nbar\n...\n"
    expected = "foo\u2603\nbar"
    codeflash_output = remove_trailing_yaml(input_str) # 1.26μs -> 1.07μs (17.2% faster)

def test_remove_trailing_yaml_marker_case_sensitive():
    # Edge: Marker is case sensitive
    input_str = "foo\nbar\n... \n"
    expected = "foo\nbar\n... \n"
    codeflash_output = remove_trailing_yaml(input_str) # 450ns -> 433ns (3.93% faster)

def test_remove_trailing_yaml_marker_with_tabs():
    # Edge: Marker with tabs (should not be removed)
    input_str = "foo\nbar\n...\n\t"
    expected = "foo\nbar\n...\n\t"
    codeflash_output = remove_trailing_yaml(input_str) # 470ns -> 457ns (2.84% faster)

# -------------------------
# 3. Large Scale Test Cases
# -------------------------

def test_remove_trailing_yaml_large_no_marker():
    # Large: Large input with no marker
    input_str = "line\n" * 1000
    expected = "line\n" * 1000
    codeflash_output = remove_trailing_yaml(input_str) # 459ns -> 451ns (1.77% faster)

def test_remove_trailing_yaml_large_with_marker():
    # Large: Large input ending with marker
    input_str = ("line\n" * 999) + "\n...\n"
    expected = ("line\n" * 999)
    codeflash_output = remove_trailing_yaml(input_str) # 1.14μs -> 937ns (21.9% faster)

def test_remove_trailing_yaml_large_marker_in_middle():
    # Large: Marker in middle of large input
    input_str = ("line\n" * 500) + "\n...\n" + ("line\n" * 499)
    expected = ("line\n" * 500) + "\n...\n" + ("line\n" * 499)
    codeflash_output = remove_trailing_yaml(input_str) # 471ns -> 433ns (8.78% faster)

def test_remove_trailing_yaml_large_multiple_markers_at_end():
    # Large: Multiple markers at end of large input
    input_str = ("line\n" * 995) + "\n...\n" * 5
    expected = ("line\n" * 995) + "\n...\n" * 4
    codeflash_output = remove_trailing_yaml(input_str) # 1.06μs -> 842ns (25.4% faster)

def test_remove_trailing_yaml_large_marker_and_extra_characters():
    # Large: Marker at end but with extra chars after
    input_str = ("line\n" * 999) + "\n...\nabc"
    expected = ("line\n" * 999) + "\n...\nabc"
    codeflash_output = remove_trailing_yaml(input_str) # 494ns -> 438ns (12.8% faster)

def test_remove_trailing_yaml_large_exact_marker():
    # Large: Input is exactly the marker, repeated
    input_str = "\n...\n"
    expected = ""
    codeflash_output = remove_trailing_yaml(input_str) # 886ns -> 718ns (23.4% faster)

def test_remove_trailing_yaml_large_marker_with_newline_before():
    # Large: Marker with extra newline before it
    input_str = ("line\n" * 999) + "\n\n...\n"
    expected = ("line\n" * 999) + "\n\n..."
    codeflash_output = remove_trailing_yaml(input_str) # 1.01μs -> 865ns (17.3% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from higgsfield.internal.experiment.params import remove_trailing_yaml

def test_remove_trailing_yaml():
    remove_trailing_yaml('\n...\n')

def test_remove_trailing_yaml_2():
    remove_trailing_yaml('')

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_0kzwu12q/tmpo9hw74br/test_concolic_coverage.py::test_remove_trailing_yaml	898ns	780ns	15.1%✅
codeflash_concolic_0kzwu12q/tmpo9hw74br/test_concolic_coverage.py::test_remove_trailing_yaml_2	505ns	485ns	4.12%✅

To edit these changes git checkout codeflash/optimize-remove_trailing_yaml-mglotke3 and push.