⚡️ Speed up function axis_spanning_shape_annotation by 66%

[codeflash-ai]

📄 66% (0.66x) speedup for axis_spanning_shape_annotation in plotly/shapeannotation.py

⏱️ Runtime : 2.10 milliseconds → 1.27 milliseconds (best of 272 runs)

📝 Explanation and details

The optimized code achieves a 65% speedup through several key performance improvements:

Data Structure Optimizations:

• Replaced lists X = [x0, x1] and Y = [y0, y1] with tuples X = (x0, x1) and Y = (y0, y1) to eliminate unnecessary list object creation overhead
• Used set literals {"top", "left"} instead of set(["top", "left"]) to avoid list creation and conversion

Eliminated Function Call Overhead:

• Replaced expensive _mean(), _argmax(), and _argmin() function calls with inline calculations:
  • _mean([y0, y1]) → (y0 + y1) / 2.0 (eliminates list creation + function call)
  • _argmax(Y) → 0 if y0 > y1 else 1 (eliminates iteration overhead)
  • max(Y) and min(Y) → direct comparisons like y0 if y0 > y1 else y1

Precomputed Common Values:

• In annotation_params_for_rect(), computed minx, maxx, miny, maxy, meanx, meany once and reused them across multiple conditions, avoiding redundant min([x0, x1]) and max([x0, x1]) calls

Optimized Dictionary Iteration:

• Used for k, v in shape_dict.items() instead of for k in shape_dict.keys() followed by shape_dict[k] lookup to eliminate repeated key lookups

Streamlined Filtering Logic:

• Replaced list(filter(lambda k: k.startswith(prefix), kwargs.keys())) with a single-pass loop that both filters keys and extracts the annotation position value, reducing iterations and lambda overhead

The optimizations are particularly effective for test cases with complex position calculations (like line positioning with "top right", "bottom left") where the function call elimination provides the biggest gains. Simple cases like returning None early also benefit significantly from the streamlined filtering logic.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from plotly.shapeannotation import axis_spanning_shape_annotation

# ========== UNIT TESTS ==========

# 1. BASIC TEST CASES

def test_vline_top_right_basic():
    # Test vline with top right position, basic usage
    shape_args = {"x0": 1, "x1": 1, "y0": 0, "y1": 10}
    kwargs = {"annotation_position": "top right"}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); result = codeflash_output # 9.45μs -> 5.61μs (68.3% faster)

def test_hline_bottom_left_basic():
    # Test hline with bottom left position, basic usage
    shape_args = {"x0": 2, "x1": 8, "y0": 5, "y1": 5}
    kwargs = {"annotation_position": "bottom left"}
    codeflash_output = axis_spanning_shape_annotation(None, "hline", shape_args, kwargs); result = codeflash_output # 7.85μs -> 4.80μs (63.6% faster)

def test_vrect_inside_top_right_basic():
    # Test vrect with inside top right position
    shape_args = {"x0": 0, "x1": 10, "y0": 3, "y1": 7}
    kwargs = {"annotation_position": "inside top right"}
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 5.86μs -> 5.36μs (9.39% faster)

def test_hrect_inside_bottom_left_basic():
    # Test hrect with inside bottom left position
    shape_args = {"x0": 4, "x1": 8, "y0": 1, "y1": 9}
    kwargs = {"annotation_position": "inside bottom left"}
    codeflash_output = axis_spanning_shape_annotation(None, "hrect", shape_args, kwargs); result = codeflash_output # 5.88μs -> 5.29μs (11.1% faster)

def test_annotation_dict_preserves_fields():
    # If annotation dict is passed, fields should be preserved unless overridden
    shape_args = {"x0": 0, "x1": 1, "y0": 2, "y1": 3}
    kwargs = {"annotation_position": "top"}
    annotation = {"text": "hello", "x": None, "y": None}
    codeflash_output = axis_spanning_shape_annotation(annotation, "vline", shape_args, kwargs); result = codeflash_output # 8.36μs -> 5.12μs (63.3% faster)

def test_annotation_kwargs_override():
    # annotation_ prefixed kwargs override annotation dict fields
    shape_args = {"x0": 0, "x1": 1, "y0": 2, "y1": 3}
    kwargs = {
        "annotation_position": "top",
        "annotation_text": "overridden",
        "annotation_font_size": 18,
    }
    annotation = {"text": "hello", "font_size": 12}
    codeflash_output = axis_spanning_shape_annotation(annotation, "vline", shape_args, kwargs); result = codeflash_output # 8.54μs -> 5.67μs (50.5% faster)

def test_default_position_is_top_right():
    # If position is None, default should be top right
    shape_args = {"x0": 0, "x1": 0, "y0": 0, "y1": 5}
    kwargs = {}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); result = codeflash_output # 1.42μs -> 663ns (114% faster)

# 2. EDGE TEST CASES

def test_none_annotation_and_no_kwargs_returns_none():
    # If annotation is None and no annotation_ kwargs, should return None
    shape_args = {"x0": 1, "x1": 2, "y0": 3, "y1": 4}
    kwargs = {}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); result = codeflash_output # 1.38μs -> 659ns (109% faster)

    # But if shape_args is empty and no kwargs, should still return None
    codeflash_output = axis_spanning_shape_annotation(None, "vline", {}, {}); result = codeflash_output # 549ns -> 228ns (141% faster)

def test_invalid_position_raises():
    # Invalid position string should raise ValueError
    shape_args = {"x0": 1, "x1": 2, "y0": 3, "y1": 4}
    kwargs = {"annotation_position": "foobar"}
    with pytest.raises(ValueError):
        axis_spanning_shape_annotation(None, "vline", shape_args, kwargs) # 8.26μs -> 5.01μs (64.8% faster)


def test_missing_shape_args_keys():
    # Missing keys in shape_args should raise KeyError
    shape_args = {"x0": 1, "y0": 2}  # missing x1, y1
    kwargs = {"annotation_position": "top right"}
    with pytest.raises(KeyError):
        axis_spanning_shape_annotation(None, "vline", shape_args, kwargs) # 3.46μs -> 2.36μs (46.4% faster)

def test_rect_inside_center():
    # Test inside center for rect
    shape_args = {"x0": 0, "x1": 10, "y0": 0, "y1": 10}
    kwargs = {"annotation_position": "inside"}
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 8.49μs -> 6.95μs (22.2% faster)

def test_rect_outside_top_left_vrect():
    # vrect, outside top left should use right/bottom anchors
    shape_args = {"x0": 2, "x1": 8, "y0": 4, "y1": 12}
    kwargs = {"annotation_position": "outside top left"}
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 7.96μs -> 6.73μs (18.2% faster)

def test_rect_outside_top_left_hrect():
    # hrect, outside top left should use left/top anchors
    shape_args = {"x0": 2, "x1": 8, "y0": 4, "y1": 12}
    kwargs = {"annotation_position": "outside top left"}
    codeflash_output = axis_spanning_shape_annotation(None, "hrect", shape_args, kwargs); result = codeflash_output # 7.00μs -> 6.14μs (13.9% faster)

def test_rect_outside_bottom_right_vrect():
    # vrect, outside bottom right should use left/top anchors
    shape_args = {"x0": 2, "x1": 8, "y0": 4, "y1": 12}
    kwargs = {"annotation_position": "outside bottom right"}
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 7.78μs -> 6.60μs (17.9% faster)

def test_rect_outside_bottom_right_hrect():
    # hrect, outside bottom right should use right/bottom anchors
    shape_args = {"x0": 2, "x1": 8, "y0": 4, "y1": 12}
    kwargs = {"annotation_position": "outside bottom right"}
    codeflash_output = axis_spanning_shape_annotation(None, "hrect", shape_args, kwargs); result = codeflash_output # 7.19μs -> 6.11μs (17.6% faster)

def test_line_left_and_right_positions():
    # vline left/right positions
    shape_args = {"x0": 5, "x1": 5, "y0": 0, "y1": 10}
    kwargs_left = {"annotation_position": "left"}
    kwargs_right = {"annotation_position": "right"}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs_left); result_left = codeflash_output # 9.55μs -> 5.31μs (79.7% faster)
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs_right); result_right = codeflash_output # 4.92μs -> 3.10μs (58.9% faster)

# 3. LARGE SCALE TEST CASES

def test_large_scale_vrect_inside_center():
    # Test large scale vrect with inside center position
    # Use 1000 for max size
    shape_args = {"x0": 0, "x1": 999, "y0": 0, "y1": 999}
    kwargs = {"annotation_position": "inside"}
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 6.46μs -> 5.29μs (22.1% faster)

def test_large_scale_hline_top_right():
    # Test large scale hline with top right position
    shape_args = {"x0": 0, "x1": 999, "y0": 500, "y1": 500}
    kwargs = {"annotation_position": "top right"}
    codeflash_output = axis_spanning_shape_annotation(None, "hline", shape_args, kwargs); result = codeflash_output # 8.05μs -> 5.21μs (54.5% faster)

def test_large_scale_vline_bottom_left():
    # Test large scale vline with bottom left position
    shape_args = {"x0": 500, "x1": 500, "y0": 0, "y1": 999}
    kwargs = {"annotation_position": "bottom left"}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); result = codeflash_output # 7.89μs -> 5.01μs (57.4% faster)

def test_large_scale_hrect_outside_top():
    # Test large scale hrect with outside top position
    shape_args = {"x0": 0, "x1": 999, "y0": 0, "y1": 999}
    kwargs = {"annotation_position": "outside top"}
    codeflash_output = axis_spanning_shape_annotation(None, "hrect", shape_args, kwargs); result = codeflash_output # 7.52μs -> 6.00μs (25.3% faster)

def test_large_scale_annotation_kwargs():
    # Large scale with annotation_ kwargs
    shape_args = {"x0": 0, "x1": 999, "y0": 0, "y1": 999}
    kwargs = {
        "annotation_position": "inside",
        "annotation_text": "big",
        "annotation_font_size": 40,
    }
    codeflash_output = axis_spanning_shape_annotation(None, "vrect", shape_args, kwargs); result = codeflash_output # 7.37μs -> 6.31μs (16.7% faster)

# 4. DETERMINISM TEST

def test_determinism_multiple_calls_same_input():
    # Multiple calls with same input should yield same output
    shape_args = {"x0": 7, "x1": 13, "y0": 2, "y1": 8}
    kwargs = {"annotation_position": "top right"}
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); out1 = codeflash_output # 7.89μs -> 4.86μs (62.3% faster)
    codeflash_output = axis_spanning_shape_annotation(None, "vline", shape_args, kwargs); out2 = codeflash_output # 4.24μs -> 2.64μs (60.7% faster)

# 5. PROPERTY PRECEDENCE TESTS

def test_property_precedence_annotation_overrides_shape():
    # If annotation dict has x/y/xanchor/yanchor, they are not overwritten
    shape_args = {"x0": 0, "x1": 10, "y0": 0, "y1": 10}
    kwargs = {"annotation_position": "inside"}
    annotation = {"x": 42, "y": 43, "xanchor": "banana", "yanchor": "apple"}
    codeflash_output = axis_spanning_shape_annotation(annotation, "vrect", shape_args, kwargs); result = codeflash_output # 6.25μs -> 5.20μs (20.1% faster)

def test_property_precedence_none_fields_are_overwritten():
    # If annotation dict has x/y/xanchor/yanchor as None, they are overwritten
    shape_args = {"x0": 0, "x1": 10, "y0": 0, "y1": 10}
    kwargs = {"annotation_position": "inside"}
    annotation = {"x": None, "y": None, "xanchor": None, "yanchor": None}
    codeflash_output = axis_spanning_shape_annotation(annotation, "vrect", shape_args, kwargs); result = codeflash_output # 6.12μs -> 4.90μs (24.8% 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 plotly.shapeannotation import axis_spanning_shape_annotation

# unit tests

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

def test_none_annotation_and_no_kwargs_returns_none():
    # If both annotation and annotation_* kwargs are missing, returns None
    codeflash_output = axis_spanning_shape_annotation(None, "vline", {"x0": 1, "x1": 1, "y0": 0, "y1": 2}, {}) # 1.33μs -> 669ns (98.8% faster)

def test_vline_top_right_default():
    # Basic vline, top right (default)
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vline",
        {"x0": 2, "x1": 2, "y0": 0, "y1": 4},
        {}
    ); result = codeflash_output # 1.31μs -> 646ns (103% faster)

def test_hline_bottom_left():
    # hline with annotation_position="bottom left"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "hline",
        {"x0": 1, "x1": 3, "y0": 5, "y1": 5},
        {"annotation_position": "bottom left"}
    ); result = codeflash_output # 9.31μs -> 5.61μs (65.9% faster)

def test_vrect_inside_top_left():
    # vrect with annotation_position="inside top left"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vrect",
        {"x0": 0, "x1": 2, "y0": 4, "y1": 8},
        {"annotation_position": "inside top left"}
    ); result = codeflash_output # 5.60μs -> 5.15μs (8.76% faster)

def test_hrect_inside_bottom_right():
    # hrect with annotation_position="inside bottom right"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "hrect",
        {"x0": 2, "x1": 6, "y0": 1, "y1": 3},
        {"annotation_position": "inside bottom right"}
    ); result = codeflash_output # 6.45μs -> 5.67μs (13.7% faster)

def test_kwargs_override_annotation_fields():
    # annotation_* kwargs override annotation fields
    codeflash_output = axis_spanning_shape_annotation(
        {"text": "foo", "x": None},  # x is None, so should be set
        "vline",
        {"x0": 3, "x1": 3, "y0": 0, "y1": 6},
        {"annotation_text": "bar", "annotation_font_size": 20}
    ); result = codeflash_output # 9.50μs -> 6.11μs (55.5% faster)

def test_existing_annotation_fields_not_overwritten():
    # Fields already set in annotation are not overwritten
    codeflash_output = axis_spanning_shape_annotation(
        {"x": 99, "xanchor": "custom", "y": None, "yanchor": None},
        "vline",
        {"x0": 4, "x1": 4, "y0": 0, "y1": 10},
        {}
    ); result = codeflash_output # 7.66μs -> 4.47μs (71.5% faster)

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


def test_invalid_position_raises():
    # If annotation_position is not valid for the shape, should raise ValueError
    with pytest.raises(ValueError):
        axis_spanning_shape_annotation(
            None,
            "vline",
            {"x0": 1, "x1": 1, "y0": 0, "y1": 2},
            {"annotation_position": "middle nowhere"}
        ) # 11.0μs -> 6.23μs (76.5% faster)

def test_rect_with_x0_equals_x1_and_y0_equals_y1():
    # Degenerate rectangle (point)
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vrect",
        {"x0": 5, "x1": 5, "y0": 7, "y1": 7},
        {"annotation_position": "inside"}
    ); result = codeflash_output # 7.51μs -> 6.22μs (20.6% faster)

def test_line_with_x0_greater_than_x1_and_y0_greater_than_y1():
    # vline with reversed coordinates
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vline",
        {"x0": 8, "x1": 8, "y0": 10, "y1": 2},
        {"annotation_position": "bottom"}
    ); result = codeflash_output # 8.79μs -> 5.32μs (65.1% faster)

def test_hrect_outside_top_right():
    # hrect, outside top right
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "hrect",
        {"x0": 1, "x1": 4, "y0": 2, "y1": 5},
        {"annotation_position": "outside top right"}
    ); result = codeflash_output # 7.49μs -> 6.67μs (12.3% faster)

def test_vrect_outside_bottom_left():
    # vrect, outside bottom left
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vrect",
        {"x0": 2, "x1": 8, "y0": 3, "y1": 9},
        {"annotation_position": "outside bottom left"}
    ); result = codeflash_output # 7.39μs -> 6.38μs (15.7% faster)

def test_hline_left_position():
    # hline with annotation_position="left"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "hline",
        {"x0": 0, "x1": 10, "y0": 7, "y1": 7},
        {"annotation_position": "left"}
    ); result = codeflash_output # 8.90μs -> 5.35μs (66.2% faster)

def test_vline_right_position():
    # vline with annotation_position="right"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vline",
        {"x0": 3, "x1": 3, "y0": 1, "y1": 9},
        {"annotation_position": "right"}
    ); result = codeflash_output # 8.12μs -> 5.07μs (60.1% faster)

def test_annotation_with_none_fields():
    # If annotation fields are None, they should be set
    codeflash_output = axis_spanning_shape_annotation(
        {"x": None, "y": None, "xanchor": None, "yanchor": None},
        "hline",
        {"x0": 2, "x1": 8, "y0": 5, "y1": 5},
        {}
    ); result = codeflash_output # 7.97μs -> 4.85μs (64.2% faster)

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

def test_many_calls_consistency():
    # Test that repeated calls with a variety of positions are consistent
    positions = [
        "top left", "top right", "top", "bottom left", "bottom right", "bottom",
        "left", "right"
    ]
    for pos in positions:
        codeflash_output = axis_spanning_shape_annotation(
            None, "vline", {"x0": 7, "x1": 7, "y0": 1, "y1": 11}, {"annotation_position": pos}
        ); res_vline = codeflash_output # 31.6μs -> 19.3μs (64.2% faster)
        codeflash_output = axis_spanning_shape_annotation(
            None, "hline", {"x0": 1, "x1": 13, "y0": 4, "y1": 4}, {"annotation_position": pos}
        ); res_hline = codeflash_output

def test_large_number_of_annotations():
    # Simulate creating a large number of annotations (but <1000 for test speed)
    for i in range(100):
        codeflash_output = axis_spanning_shape_annotation(
            None,
            "hrect" if i % 2 == 0 else "vrect",
            {"x0": i, "x1": i + 10, "y0": i * 2, "y1": i * 2 + 5},
            {"annotation_position": "inside bottom right"}
        ); res = codeflash_output # 219μs -> 176μs (24.5% faster)
        if i % 2 == 0:
            pass
        else:
            pass

def test_large_rect_coordinates():
    # Test with large coordinate values
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "vrect",
        {"x0": 1_000_000, "x1": 2_000_000, "y0": -1_000_000, "y1": 3_000_000},
        {"annotation_position": "inside top right"}
    ); result = codeflash_output # 5.37μs -> 4.82μs (11.5% faster)

def test_large_kwargs():
    # Test with many annotation_* kwargs (simulate real-world usage)
    kwargs = {f"annotation_custom{i}": i for i in range(50)}
    kwargs["annotation_position"] = "inside"
    codeflash_output = axis_spanning_shape_annotation(
        None,
        "hrect",
        {"x0": 0, "x1": 10, "y0": 0, "y1": 10},
        kwargs
    ); result = codeflash_output # 16.2μs -> 15.1μs (7.41% faster)
    for i in range(50):
        pass

def test_performance_large_loop():
    # Ensure function does not degrade with many calls (not a strict perf test, but checks no errors)
    for i in range(500):
        axis_spanning_shape_annotation(
            None,
            "hline" if i % 2 == 0 else "vline",
            {"x0": i, "x1": i + 1, "y0": i * 2, "y1": i * 2 + 1},
            {"annotation_position": "bottom right"}
        ) # 1.51ms -> 834μs (81.6% 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-axis_spanning_shape_annotation-mhgeceng and push.