⚡️ Speed up function setp by 62%

[codeflash-ai]

📄 62% (0.62x) speedup for setp in lib/matplotlib/pyplot.py

⏱️ Runtime : 7.48 milliseconds → 4.62 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves a 61% speedup by replacing the recursive cbook.flatten calls with a custom _flatten_fast function that uses an iterative stack-based approach.

Key optimization changes:

1. Custom iterative flattening: The new _flatten_fast function eliminates recursion overhead by using a stack-based iteration. This avoids the function call overhead and stack frame creation that comes with the recursive cbook.flatten implementation.

2. Two critical flatten replacements:

   • objs = list(cbook.flatten(obj)) → objs = list(_flatten_fast(obj))
   • return list(cbook.flatten(ret)) → return list(_flatten_fast(ret))

Why this optimization works:

• Reduced call stack depth: The iterative approach eliminates recursive function calls, which are expensive in Python due to frame creation and teardown
• Better memory access patterns: Stack-based iteration has more predictable memory access compared to recursion
• Lower overhead: Fewer function calls mean less bytecode interpretation overhead

Impact on workloads:

The function reference shows setp is called from matplotlib.pyplot, making it part of the public plotting API. The test results demonstrate the optimization is particularly effective for:

• Large artist collections: 73-136% faster for 500-999 artists
• Nested structures: 7-9% faster for nested artist lists
• Batch operations: 76% faster for 1000-artist batches

This optimization significantly benefits matplotlib users who work with many plot objects simultaneously, such as creating complex visualizations with hundreds of lines, points, or other artists that need property updates.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_artist.py::test_setp	348μs	356μs	-2.23%⚠️
test_polar.py::test_polar_gridlines	315μs	308μs	2.11%✅
test_table.py::test_table_cells	368μs	369μs	-0.233%⚠️
test_ticker.py::test_remove_overlap	499μs	485μs	2.84%✅

🌀 Generated Regression Tests and Runtime

# imports
from matplotlib.pyplot import setp

# --- Minimal stubs for testing setp, Artist, and cbook.flatten ---


class Artist:
    """
    Minimal Artist stub for testing.
    Implements set() and update() methods for property setting.
    """

    def __init__(self):
        self.props = {}
        self.set_calls = []
        self.update_calls = []

    # Accepts only known properties for testing
    _valid_properties = {
        "color": {"r", "g", "b", "c", "m", "y", "k", "w", "red", "green", "blue"},
        "linewidth": set(range(1, 11)),
        "linestyle": {"-", "--", "-.", ":", "", "solid", "dashed"},
        "alpha": set([None] + [x / 10 for x in range(11)]),
    }

    def set(self, **kwargs):
        # Record the call for test verification
        self.set_calls.append(kwargs.copy())
        for k, v in kwargs.items():
            if k not in self._valid_properties:
                raise AttributeError(f"Unknown property: {k}")
            # Simulate allowed values
            if v not in self._valid_properties[k]:
                raise ValueError(f"Invalid value {v} for property {k}")
            self.props[k] = v
        return self

    def update(self, d):
        # Record the call for test verification
        self.update_calls.append(d.copy())
        for k, v in d.items():
            if k not in self._valid_properties:
                raise AttributeError(f"Unknown property: {k}")
            if v not in self._valid_properties[k]:
                raise ValueError(f"Invalid value {v} for property {k}")
            self.props[k] = v
        return self


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

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


def test_set_multiple_artists():
    # Test setting properties on multiple Artists at once
    a1 = Artist()
    a2 = Artist()
    setp([a1, a2], color="k", linewidth=2)  # 35.8μs -> 33.7μs (6.27% faster)


def test_setp_returns_list_of_artists():
    # Test that setp returns a list of the objects (flattened)
    a1 = Artist()
    a2 = Artist()
    codeflash_output = setp([a1, a2], color="r")
    result = codeflash_output  # 29.6μs -> 28.6μs (3.34% faster)


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


def test_empty_iterable_returns_none():
    # Test that an empty iterable returns None and does not raise
    codeflash_output = setp([], color="r")
    result = codeflash_output  # 4.56μs -> 5.68μs (19.7% slower)


def test_nested_iterable_of_artists():
    # Test that nested iterables of Artists are flattened and all set
    a1 = Artist()
    a2 = Artist()
    a3 = Artist()
    setp([[a1, [a2]], a3], color="g")  # 36.6μs -> 35.6μs (2.83% faster)


def test_setp_many_artists():
    # Test setting a property on a large number of Artists
    artists = [Artist() for _ in range(500)]
    setp(artists, color="b", linewidth=5)  # 946μs -> 546μs (73.4% faster)
    for a in artists:
        pass


def test_setp_deeply_nested_artists():
    # Test flattening a deeply nested structure of Artists
    artists = [[Artist(), [Artist(), [Artist()]]], Artist()]
    setp(artists, color="c")  # 37.8μs -> 34.7μs (8.99% faster)
    # Flatten for checking
    flat = []

    def _flat(seq):
        for x in seq:
            if isinstance(x, Artist):
                flat.append(x)
            else:
                _flat(x)

    _flat(artists)
    for a in flat:
        pass


def test_setp_performance_large_batch(monkeypatch):
    # Test that setp does not take excessive time for 1000 artists
    # (Performance test, but also checks correctness)
    artists = [Artist() for _ in range(1000)]
    setp(artists, color="r", linewidth=2)  # 1.80ms -> 1.02ms (76.7% faster)
    for a in artists:
        pass


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

# imports
import pytest
from matplotlib.pyplot import setp

# Minimal stubs for matplotlib internals used by setp, so tests are self-contained
# (We do NOT mock setp or its logic, but we must provide minimal Artist etc.)


class DummyArtist:
    """A minimal Artist stub for testing setp."""

    # Simulate settable properties and their allowed values
    _settable = {
        "color": {"r", "g", "b", "k", "w"},
        "linewidth": {1, 2, 3, 4, 5},
        "linestyle": {"-", "--", "-.", ":", ""},
        "alpha": {None, 0.0, 0.5, 1.0},
    }

    def __init__(self):
        self.props = {}
        self.updated = False
        self.set_called = False

    def update(self, d):
        self.props.update(d)
        self.updated = True
        return True

    def set(self, **kwargs):
        self.props.update(kwargs)
        self.set_called = True
        return True

    @classmethod
    def get_setters(cls):
        return cls._settable.keys()

    @classmethod
    def get_valid_values(cls, prop):
        return cls._settable.get(prop, set())


# --------------------------
# UNIT TESTS FOR setp
# --------------------------

# 1. BASIC TEST CASES


def test_set_properties_on_multiple_artists():
    # Test setting properties on a list of artists
    arts = [DummyArtist(), DummyArtist()]
    setp(arts, color="k", linewidth=5)  # 35.6μs -> 36.8μs (3.24% slower)
    for art in arts:
        pass


def test_empty_artist_list_returns_none():
    # Test that setp returns None for an empty list
    codeflash_output = setp([])
    result = codeflash_output  # 4.11μs -> 5.07μs (18.9% slower)


def test_non_iterable_object_raises():
    # Test that passing a non-Artist, non-iterable object raises TypeError
    with pytest.raises(TypeError):
        setp(42, color="r")  # 4.45μs -> 5.06μs (12.0% slower)


def test_setp_with_nested_artist_list():
    # Test nested list of artists is handled and all are updated
    arts = [[DummyArtist(), DummyArtist()], [DummyArtist()]]
    setp(arts, color="w")  # 39.2μs -> 36.5μs (7.55% faster)
    for sublist in arts:
        for art in sublist:
            pass


def test_setp_large_artist_list():
    # Test setp with a large list of artists
    arts = [DummyArtist() for _ in range(500)]
    setp(arts, color="g", linewidth=4)  # 778μs -> 356μs (119% faster)
    for art in arts:
        pass


def test_setp_large_nested_artist_list():
    # Test setp with a large, deeply nested list of artists
    nested_arts = [[DummyArtist() for _ in range(10)] for _ in range(50)]
    setp(nested_arts, color="b")  # 762μs -> 337μs (126% faster)
    for sublist in nested_arts:
        for art in sublist:
            pass


def test_setp_performance_large_scale():
    # Test that setp does not become quadratic for large input (smoke test)
    arts = [DummyArtist() for _ in range(999)]
    setp(arts, color="k")  # 1.41ms -> 595μs (136% faster)
    # If it completes, it's a pass


# Edge: test that setp returns a list of results for property setting
def test_setp_returns_list_of_results():
    arts = [DummyArtist(), DummyArtist()]
    codeflash_output = setp(arts, color="r")
    result = codeflash_output  # 30.9μs -> 29.2μs (5.77% faster)


# Edge: test that setp prints to sys.stdout by default

To edit these changes git checkout codeflash/optimize-setp-mjacw83e and push.