⚡️ Speed up method CallbackRegistry._remove_proxy by 49%

[codeflash-ai]

📄 49% (0.49x) speedup for CallbackRegistry._remove_proxy in lib/matplotlib/cbook.py

⏱️ Runtime : 1.76 milliseconds → 1.18 milliseconds (best of 36 runs)

📝 Explanation and details

The optimization achieves a 49% speedup by eliminating redundant operations and reducing attribute lookups in the hot loop. Here are the key improvements:

1. Eliminated list() conversion overhead:

• Original: list(self._func_cid_map.items()) creates an unnecessary copy of all items
• Optimized: Direct iteration over func_cid_map.items() saves memory allocation and copying

2. Reduced attribute lookups:

• Original: Repeated self.callbacks[signal], self._func_cid_map, self._pickled_cids lookups in the loop
• Optimized: Pre-cached these attributes as local variables (callbacks, func_cid_map, pickled_cids)

3. Avoided wasteful .pop() calls:

• Original: Called proxy_to_cid.pop(proxy, None) on every iteration, even when proxy wasn't present
• Optimized: Added if proxy in proxy_to_cid: check first, then called proxy_to_cid.pop(proxy) only when needed

4. Simplified callback removal:

• Original: Used del self.callbacks[signal][cid]
• Optimized: Used callbacks[signal].pop(cid, None) which is more defensive

The line profiler shows the optimization's impact:

• Loop iteration time reduced from 2.9ms to 2.8ms (5% improvement)
• .pop() overhead reduced from 4.8ms to 4.4ms (8% improvement) due to fewer unnecessary calls
• Overall function time reduced from 11.8ms to 9.5ms

Performance characteristics:

• Best case: When proxies are not found (32.7% speedup in tests) - avoids all .pop() calls
• Typical case: With mixed found/not-found proxies (18-27% speedup in most tests)
• Large scale: Maintains performance gains even with 500+ callbacks (21% speedup)

This optimization is particularly valuable since _remove_proxy is called during cleanup operations when callbacks are disconnected or when weak references are garbage collected, making it performance-critical for applications with many dynamic callbacks.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# imports
import pytest
from matplotlib.cbook import CallbackRegistry


# function to test (copied from above, with minimal dependencies)
def _exception_printer(exc):
    pass  # not used in these tests


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


# Helper: create a dummy proxy object
class DummyProxy:
    def __init__(self, value):
        self.value = value

    def __call__(self):
        return self.value

    def __eq__(self, other):
        return isinstance(other, DummyProxy) and self.value == other.value

    def __hash__(self):
        return hash(self.value)


@pytest.fixture
def registry():
    return CallbackRegistry()


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


def test_remove_proxy_removes_entry_basic(registry):
    """Test that _remove_proxy removes the correct callback and mapping for a single entry."""
    signal = "test"
    cid = 42
    proxy = DummyProxy("func")
    registry.callbacks = {signal: {cid: proxy}}
    registry._func_cid_map = {signal: {proxy: cid}}
    registry._pickled_cids = {cid}
    registry._remove_proxy(proxy)  # 3.62μs -> 3.05μs (18.8% faster)


def test_remove_proxy_does_not_affect_other_signals(registry):
    """Test that _remove_proxy only removes from the relevant signal."""
    sig1, sig2 = "a", "b"
    cid1, cid2 = 1, 2
    proxy1 = DummyProxy("f1")
    proxy2 = DummyProxy("f2")
    registry.callbacks = {sig1: {cid1: proxy1}, sig2: {cid2: proxy2}}
    registry._func_cid_map = {sig1: {proxy1: cid1}, sig2: {proxy2: cid2}}
    registry._pickled_cids = {cid1, cid2}
    registry._remove_proxy(proxy1)  # 3.60μs -> 3.01μs (19.8% faster)


def test_remove_proxy_removes_only_one_proxy(registry):
    """Test that _remove_proxy only removes the specified proxy, not others."""
    signal = "sig"
    cid1, cid2 = 1, 2
    proxy1 = DummyProxy("f1")
    proxy2 = DummyProxy("f2")
    registry.callbacks = {signal: {cid1: proxy1, cid2: proxy2}}
    registry._func_cid_map = {signal: {proxy1: cid1, proxy2: cid2}}
    registry._pickled_cids = {cid1, cid2}
    registry._remove_proxy(proxy1)  # 3.22μs -> 2.74μs (17.4% faster)


def test_remove_proxy_returns_if_proxy_not_found(registry):
    """Test that _remove_proxy is a no-op if proxy is not present."""
    signal = "sig"
    cid = 1
    proxy = DummyProxy("f1")
    registry.callbacks = {signal: {cid: proxy}}
    registry._func_cid_map = {signal: {proxy: cid}}
    registry._pickled_cids = {cid}
    missing_proxy = DummyProxy("notfound")
    # Should not raise
    registry._remove_proxy(missing_proxy)  # 2.55μs -> 1.92μs (32.7% faster)


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


def test_remove_proxy_on_empty_registry(registry):
    """Test that _remove_proxy does nothing on empty registry."""
    proxy = DummyProxy("x")
    registry._remove_proxy(proxy)  # 1.86μs -> 1.47μs (26.5% faster)


def test_remove_proxy_when_signal_dict_becomes_empty(registry):
    """Test that _remove_proxy deletes signal and func_cid_map when last proxy is removed."""
    signal = "sig"
    cid = 100
    proxy = DummyProxy("func")
    registry.callbacks = {signal: {cid: proxy}}
    registry._func_cid_map = {signal: {proxy: cid}}
    registry._pickled_cids = {cid}
    registry._remove_proxy(proxy)  # 3.53μs -> 3.02μs (17.1% faster)


def test_remove_proxy_with_multiple_signals_and_proxies(registry):
    """Test correct removal when there are multiple signals, each with multiple proxies."""
    sigs = ["s1", "s2"]
    cids = [1, 2, 3, 4]
    proxies = [DummyProxy(f"f{i}") for i in range(4)]
    registry.callbacks = {
        sigs[0]: {cids[0]: proxies[0], cids[1]: proxies[1]},
        sigs[1]: {cids[2]: proxies[2], cids[3]: proxies[3]},
    }
    registry._func_cid_map = {
        sigs[0]: {proxies[0]: cids[0], proxies[1]: cids[1]},
        sigs[1]: {proxies[2]: cids[2], proxies[3]: cids[3]},
    }
    registry._pickled_cids = set(cids)
    # Remove proxies[1] (should leave sigs[0] with one proxy)
    registry._remove_proxy(proxies[1])  # 3.33μs -> 2.63μs (26.5% faster)
    # Remove proxies[0] (should remove sigs[0] entirely)
    registry._remove_proxy(proxies[0])  # 1.33μs -> 1.24μs (6.68% faster)


def test_remove_proxy_does_not_fail_if_sys_is_finalizing(monkeypatch, registry):
    """Test that _remove_proxy returns immediately if sys.is_finalizing() is True."""
    signal = "sig"
    cid = 1
    proxy = DummyProxy("f")
    registry.callbacks = {signal: {cid: proxy}}
    registry._func_cid_map = {signal: {proxy: cid}}
    registry._pickled_cids = {cid}
    # Monkeypatch sys.is_finalizing to return True
    registry._remove_proxy(
        proxy, _is_finalizing=lambda: True
    )  # 868ns -> 908ns (4.41% slower)


def test_remove_proxy_with_non_hashable_proxy(registry):
    """Test that _remove_proxy handles proxies that are not hashable (should not crash)."""

    class NonHashableProxy:
        def __call__(self):
            return 1

        __hash__ = None

        def __eq__(self, other):
            return False

    proxy = NonHashableProxy()
    # Place in a list, not in dict, so _remove_proxy can't find it
    registry.callbacks = {}
    registry._func_cid_map = {}
    # Should not raise
    registry._remove_proxy(proxy)  # 1.77μs -> 1.53μs (15.9% faster)


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


def test_remove_proxy_large_scale_removal():
    """Test _remove_proxy performance and correctness with large number of signals and proxies."""
    registry = CallbackRegistry()
    N_signals = 50
    N_proxies_per_signal = 15
    proxies = []
    for i in range(N_signals):
        signal = f"signal_{i}"
        registry.callbacks[signal] = {}
        registry._func_cid_map[signal] = {}
        for j in range(N_proxies_per_signal):
            cid = i * N_proxies_per_signal + j
            proxy = DummyProxy(f"f_{i}_{j}")
            proxies.append((signal, cid, proxy))
            registry.callbacks[signal][cid] = proxy
            registry._func_cid_map[signal][proxy] = cid
            registry._pickled_cids.add(cid)
    # Remove every 5th proxy
    for idx, (signal, cid, proxy) in enumerate(proxies):
        if idx % 5 == 0:
            registry._remove_proxy(proxy)
    # Check that removed proxies are gone
    for idx, (signal, cid, proxy) in enumerate(proxies):
        if idx % 5 == 0:
            # Should be removed
            if signal in registry.callbacks:
                pass
            if signal in registry._func_cid_map:
                pass
        else:
            pass


def test_remove_proxy_large_scale_signal_cleanup():
    """Test that signals are deleted when all their proxies are removed in a large registry."""
    registry = CallbackRegistry()
    N_signals = 10
    N_proxies = 10
    proxies = []
    for i in range(N_signals):
        signal = f"signal_{i}"
        registry.callbacks[signal] = {}
        registry._func_cid_map[signal] = {}
        for j in range(N_proxies):
            cid = i * N_proxies + j
            proxy = DummyProxy(f"f_{i}_{j}")
            proxies.append((signal, cid, proxy))
            registry.callbacks[signal][cid] = proxy
            registry._func_cid_map[signal][proxy] = cid
            registry._pickled_cids.add(cid)
    # Remove all proxies for each signal, one signal at a time
    for i in range(N_signals):
        signal = f"signal_{i}"
        for j in range(N_proxies):
            cid = i * N_proxies + j
            proxy = DummyProxy(f"f_{i}_{j}")
            registry._remove_proxy(proxy)


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

import weakref

# imports
import pytest
from matplotlib.cbook import CallbackRegistry


# function to test (copied from above, with minimal external dependencies)
def _exception_printer(exc):
    print("Exception:", exc)


def _weak_or_strong_ref(func, remove_proxy):
    # Helper to create a weakref proxy for a bound method, or just return func
    # For testing, we simulate weakref for bound methods, strong for others.
    if hasattr(func, "__self__") and func.__self__ is not None:
        # Bound method: use weakref
        proxy = weakref.WeakMethod(func, remove_proxy)
    else:
        # Function: use strong ref
        proxy = func
    return proxy


# --- Unit tests for CallbackRegistry._remove_proxy ---


@pytest.fixture
def registry():
    # Create a fresh CallbackRegistry for each test
    return CallbackRegistry()


# Helper for creating a bound method proxy
class Dummy:
    def cb(self):
        pass


def make_bound_proxy(registry, instance=None):
    if instance is None:
        instance = Dummy()
    return _weak_or_strong_ref(instance.cb, registry._remove_proxy)


def make_func_proxy(registry):
    def cb():
        pass

    return _weak_or_strong_ref(cb, registry._remove_proxy)


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


def test_remove_proxy_removes_callback_and_func_cid_map(registry):
    # Setup: add a signal, callback, and mapping
    signal = "test"
    cid = 42
    proxy = make_func_proxy(registry)
    registry.callbacks[signal] = {cid: proxy}
    registry._func_cid_map[signal] = {proxy: cid}
    registry._pickled_cids.add(cid)
    # Action: remove proxy
    registry._remove_proxy(proxy)  # 3.42μs -> 2.74μs (24.9% faster)


def test_remove_proxy_removes_only_one_signal(registry):
    # Setup: two signals, only one should be affected
    proxy1 = make_func_proxy(registry)
    proxy2 = make_func_proxy(registry)
    registry.callbacks["sig1"] = {1: proxy1}
    registry.callbacks["sig2"] = {2: proxy2}
    registry._func_cid_map["sig1"] = {proxy1: 1}
    registry._func_cid_map["sig2"] = {proxy2: 2}
    registry._pickled_cids.update([1, 2])
    # Remove proxy1
    registry._remove_proxy(proxy1)  # 3.37μs -> 2.71μs (24.3% faster)


def test_remove_proxy_multiple_callbacks_in_signal(registry):
    # Setup: signal with multiple callbacks
    proxy1 = make_func_proxy(registry)
    proxy2 = make_func_proxy(registry)
    registry.callbacks["sig"] = {1: proxy1, 2: proxy2}
    registry._func_cid_map["sig"] = {proxy1: 1, proxy2: 2}
    registry._pickled_cids.update([1, 2])
    # Remove proxy1
    registry._remove_proxy(proxy1)  # 3.15μs -> 2.33μs (34.9% faster)


def test_remove_proxy_no_matching_proxy(registry):
    # Setup: registry with unrelated proxy
    proxy1 = make_func_proxy(registry)
    proxy2 = make_func_proxy(registry)
    registry.callbacks["sig"] = {1: proxy1}
    registry._func_cid_map["sig"] = {proxy1: 1}
    # Remove proxy2 (not present)
    registry._remove_proxy(proxy2)  # 2.30μs -> 1.67μs (37.6% faster)


def test_remove_proxy_when_signal_dict_becomes_empty(registry):
    # Setup: signal with one callback
    proxy = make_func_proxy(registry)
    registry.callbacks["sig"] = {1: proxy}
    registry._func_cid_map["sig"] = {proxy: 1}
    registry._pickled_cids.add(1)
    # Remove proxy
    registry._remove_proxy(proxy)  # 3.37μs -> 2.65μs (27.3% faster)


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


def test_remove_proxy_with_multiple_signals_and_proxies(registry):
    # Setup: multiple signals, proxies, overlapping cids
    proxy1 = make_func_proxy(registry)
    proxy2 = make_func_proxy(registry)
    proxy3 = make_func_proxy(registry)
    registry.callbacks["sig1"] = {1: proxy1, 2: proxy2}
    registry.callbacks["sig2"] = {3: proxy3}
    registry._func_cid_map["sig1"] = {proxy1: 1, proxy2: 2}
    registry._func_cid_map["sig2"] = {proxy3: 3}
    registry._pickled_cids.update([1, 2, 3])
    # Remove proxy2
    registry._remove_proxy(proxy2)  # 3.13μs -> 2.42μs (29.6% faster)


def test_remove_proxy_with_empty_func_cid_map(registry):
    # Setup: func_cid_map is empty
    registry._func_cid_map = {}
    proxy = make_func_proxy(registry)
    # Remove proxy should do nothing and not error
    registry._remove_proxy(proxy)  # 1.88μs -> 1.58μs (19.3% faster)


def test_remove_proxy_with_empty_callbacks(registry):
    # Setup: callbacks is empty
    registry.callbacks = {}
    proxy = make_func_proxy(registry)
    # Remove proxy should do nothing and not error
    registry._remove_proxy(proxy)  # 1.88μs -> 1.59μs (18.5% faster)


def test_remove_proxy_with_duplicate_proxies(registry):
    # Setup: two signals, same proxy object (simulate bug)
    proxy = make_func_proxy(registry)
    registry.callbacks["sig1"] = {1: proxy}
    registry.callbacks["sig2"] = {2: proxy}
    registry._func_cid_map["sig1"] = {proxy: 1}
    registry._func_cid_map["sig2"] = {proxy: 2}
    registry._pickled_cids.update([1, 2])
    # Remove proxy should remove only one signal (first found)
    registry._remove_proxy(proxy)  # 3.48μs -> 2.59μs (34.2% faster)
    # Only one signal removed, the other remains
    removed = []
    for sig in ["sig1", "sig2"]:
        if sig not in registry.callbacks:
            removed.append(sig)
    remaining = set(["sig1", "sig2"]) - set(removed)
    for sig in remaining:
        pass


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


def test_remove_proxy_large_number_of_signals_and_callbacks(registry):
    # Setup: 100 signals, each with 5 callbacks
    num_signals = 100
    num_callbacks = 5
    proxies = []
    for i in range(num_signals):
        signal = f"sig{i}"
        registry.callbacks[signal] = {}
        registry._func_cid_map[signal] = {}
        for j in range(num_callbacks):
            cid = i * num_callbacks + j
            proxy = make_func_proxy(registry)
            proxies.append((signal, cid, proxy))
            registry.callbacks[signal][cid] = proxy
            registry._func_cid_map[signal][proxy] = cid
            registry._pickled_cids.add(cid)
    # Remove every third proxy
    for idx, (signal, cid, proxy) in enumerate(proxies):
        if idx % 3 == 0:
            registry._remove_proxy(proxy)
    # Assert: removed proxies/cids gone, signals with no callbacks removed
    for i in range(num_signals):
        signal = f"sig{i}"
        # If all callbacks for a signal removed, signal is gone
        if all(
            idx % 3 == 0 for idx in range(i * num_callbacks, (i + 1) * num_callbacks)
        ):
            pass
        else:
            # Otherwise, only some callbacks remain
            if signal in registry.callbacks:
                for j in range(num_callbacks):
                    idx = i * num_callbacks + j
                    cid = idx
                    proxy = proxies[idx][2]
                    if idx % 3 == 0:
                        pass
                    else:
                        pass


def test_remove_proxy_performance_under_many_removals(registry):
    # Setup: 500 callbacks in one signal
    signal = "sig"
    num_callbacks = 500
    proxies = []
    registry.callbacks[signal] = {}
    registry._func_cid_map[signal] = {}
    for cid in range(num_callbacks):
        proxy = make_func_proxy(registry)
        proxies.append((cid, proxy))
        registry.callbacks[signal][cid] = proxy
        registry._func_cid_map[signal][proxy] = cid
        registry._pickled_cids.add(cid)
    # Remove all proxies
    for cid, proxy in proxies:
        registry._remove_proxy(proxy)  # 259μs -> 214μs (21.0% faster)
    for cid, _ in proxies:
        pass


def test_remove_proxy_large_scale_with_mixed_bound_and_unbound(registry):
    # Setup: 50 signals, each with 10 bound and 10 unbound callbacks
    num_signals = 50
    num_callbacks = 10
    proxies = []
    for i in range(num_signals):
        signal = f"sig{i}"
        registry.callbacks[signal] = {}
        registry._func_cid_map[signal] = {}
        # Bound methods
        for j in range(num_callbacks):
            cid = i * 2 * num_callbacks + j
            proxy = make_bound_proxy(registry)
            proxies.append((signal, cid, proxy))
            registry.callbacks[signal][cid] = proxy
            registry._func_cid_map[signal][proxy] = cid
            registry._pickled_cids.add(cid)
        # Unbound functions
        for j in range(num_callbacks):
            cid = i * 2 * num_callbacks + num_callbacks + j
            proxy = make_func_proxy(registry)
            proxies.append((signal, cid, proxy))
            registry.callbacks[signal][cid] = proxy
            registry._func_cid_map[signal][proxy] = cid
            registry._pickled_cids.add(cid)
    # Remove all bound proxies
    for idx, (signal, cid, proxy) in enumerate(proxies):
        # Bound proxies are first half
        if idx % (2 * num_callbacks) < num_callbacks:
            registry._remove_proxy(proxy)
    # Assert: all bound cids/proxies gone, unbound remain
    for i in range(num_signals):
        signal = f"sig{i}"
        for j in range(2 * num_callbacks):
            idx = i * 2 * num_callbacks + j
            cid = proxies[idx][1]
            proxy = proxies[idx][2]
            if j < num_callbacks:
                # Bound: should be gone
                if signal in registry.callbacks:
                    pass
                if signal in registry._func_cid_map:
                    pass
            else:
                pass


# 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-CallbackRegistry._remove_proxy-misaz0hr and push.