Support unhashable callbacks in CallbackRegistry.

anntzer · anntzer · commit 87f194f17785 · 2023-07-13T22:21:46.000+02:00
... by replacing _func_cid_map by a dict-like structure (_UnhashDict)
that also supports unhashable entries.

Note that _func_cid_map (and thus _UnhashDict) can be dropped if we get
rid of proxy deduplication in CallbackRegistry.
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -112,6 +112,61 @@ def _weak_or_strong_ref(func, callback):
         return _StrongRef(func)
 
 
+class _UnhashDict:
+    """
+    A minimal dict-like class that also supports unhashable keys, storing them
+    in a list of key-value pairs.
+
+    This class only implements the interface needed for `CallbackRegistry`, and
+    tries to minimize the overhead for the hashable case.
+    """
+
+    def __init__(self, pairs):
+        self._dict = {}
+        self._pairs = []
+        for k, v in pairs:
+            self[k] = v
+
+    def __setitem__(self, key, value):
+        try:
+            self._dict[key] = value
+        except TypeError:
+            for i, (k, v) in enumerate(self._pairs):
+                if k == key:
+                    self._pairs[i] = (key, value)
+                    break
+            else:
+                self._pairs.append((key, value))
+
+    def __getitem__(self, key):
+        try:
+            return self._dict[key]
+        except TypeError:
+            pass
+        for k, v in self._pairs:
+            if k == key:
+                return v
+        raise KeyError(key)
+
+    def pop(self, key, *args):
+        try:
+            if key in self._dict:
+                return self._dict.pop(key)
+        except TypeError:
+            for i, (k, v) in enumerate(self._pairs):
+                if k == key:
+                    del self._pairs[i]
+                    return v
+        if args:
+            return args[0]
+        raise KeyError(key)
+
+    def __iter__(self):
+        yield from self._dict
+        for k, v in self._pairs:
+            yield k
+
+
 class CallbackRegistry:
     """
     Handle registering, processing, blocking, and disconnecting
@@ -178,7 +233,7 @@ def __init__(self, exception_handler=_exception_printer, *, signals=None):
         self.exception_handler = exception_handler
         self.callbacks = {}
         self._cid_gen = itertools.count()
-        self._func_cid_map = {}
+        self._func_cid_map = _UnhashDict([])
         # A hidden variable that marks cids that need to be pickled.
         self._pickled_cids = set()
 
@@ -202,9 +257,9 @@ def __setstate__(self, state):
             s: {cid: _weak_or_strong_ref(func, functools.partial(self._remove_proxy, s))
                 for cid, func in d.items()}
             for s, d in self.callbacks.items()}
-        self._func_cid_map = {
-            (s, proxy): cid
-            for s, d in self.callbacks.items() for cid, proxy in d.items()}
+        self._func_cid_map = _UnhashDict(
+            ((s, proxy), cid)
+            for s, d in self.callbacks.items() for cid, proxy in d.items())
         self._cid_gen = itertools.count(cid_count)
 
     def connect(self, signal, func):
@@ -258,7 +313,7 @@ def disconnect(self, cid):
             return
         assert self.callbacks[signal][cid] == proxy
         del self.callbacks[signal][cid]
-        del self._func_cid_map[signal, proxy]
+        self._func_cid_map.pop((signal, proxy))
         if len(self.callbacks[signal]) == 0:  # Clean up empty dicts
             del self.callbacks[signal]
 
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -178,6 +178,15 @@ def test_boxplot_stats_autorange_false(self):
         assert_array_almost_equal(bstats_true[0]['fliers'], [])
 
 
+class Hashable:
+    def dummy(self): pass
+
+
+class Unhashable:
+    __hash__ = None  # type: ignore
+    def dummy(self): pass
+
+
 class Test_callback_registry:
     def setup_method(self):
         self.signal = 'test'
@@ -200,13 +209,13 @@ def count(self):
 
     def is_empty(self):
         np.testing.break_cycles()
-        assert self.callbacks._func_cid_map == {}
+        assert [*self.callbacks._func_cid_map] == []
         assert self.callbacks.callbacks == {}
         assert self.callbacks._pickled_cids == set()
 
     def is_not_empty(self):
         np.testing.break_cycles()
-        assert self.callbacks._func_cid_map != {}
+        assert [*self.callbacks._func_cid_map] != []
         assert self.callbacks.callbacks != {}
 
     def test_cid_restore(self):
@@ -217,12 +226,13 @@ def test_cid_restore(self):
         assert cid == 1
 
     @pytest.mark.parametrize('pickle', [True, False])
-    def test_callback_complete(self, pickle):
+    @pytest.mark.parametrize('cls', [Hashable, Unhashable])
+    def test_callback_complete(self, pickle, cls):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
-        mini_me = Test_callback_registry()
+        mini_me = cls()
 
         # test that we can add a callback
         cid1 = self.connect(self.signal, mini_me.dummy, pickle)
@@ -233,7 +243,7 @@ def test_callback_complete(self, pickle):
         cid2 = self.connect(self.signal, mini_me.dummy, pickle)
         assert cid1 == cid2
         self.is_not_empty()
-        assert len(self.callbacks._func_cid_map) == 1
+        assert len([*self.callbacks._func_cid_map]) == 1
         assert len(self.callbacks.callbacks) == 1
 
         del mini_me
@@ -242,12 +252,13 @@ def test_callback_complete(self, pickle):
         self.is_empty()
 
     @pytest.mark.parametrize('pickle', [True, False])
-    def test_callback_disconnect(self, pickle):
+    @pytest.mark.parametrize('cls', [Hashable, Unhashable])
+    def test_callback_disconnect(self, pickle, cls):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
-        mini_me = Test_callback_registry()
+        mini_me = cls()
 
         # test that we can add a callback
         cid1 = self.connect(self.signal, mini_me.dummy, pickle)
@@ -260,12 +271,13 @@ def test_callback_disconnect(self, pickle):
         self.is_empty()
 
     @pytest.mark.parametrize('pickle', [True, False])
-    def test_callback_wrong_disconnect(self, pickle):
+    @pytest.mark.parametrize('cls', [Hashable, Unhashable])
+    def test_callback_wrong_disconnect(self, pickle, cls):
         # ensure we start with an empty registry
         self.is_empty()
 
         # create a class for testing
-        mini_me = Test_callback_registry()
+        mini_me = cls()
 
         # test that we can add a callback
         cid1 = self.connect(self.signal, mini_me.dummy, pickle)
@@ -278,20 +290,21 @@ def test_callback_wrong_disconnect(self, pickle):
         self.is_not_empty()
 
     @pytest.mark.parametrize('pickle', [True, False])
-    def test_registration_on_non_empty_registry(self, pickle):
+    @pytest.mark.parametrize('cls', [Hashable, Unhashable])
+    def test_registration_on_non_empty_registry(self, pickle, cls):
         # ensure we start with an empty registry
         self.is_empty()
 
         # setup the registry with a callback
-        mini_me = Test_callback_registry()
+        mini_me = cls()
         self.connect(self.signal, mini_me.dummy, pickle)
 
         # Add another callback
-        mini_me2 = Test_callback_registry()
+        mini_me2 = cls()
         self.connect(self.signal, mini_me2.dummy, pickle)
 
         # Remove and add the second callback
-        mini_me2 = Test_callback_registry()
+        mini_me2 = cls()
         self.connect(self.signal, mini_me2.dummy, pickle)
 
         # We still have 2 references
@@ -303,9 +316,6 @@ def test_registration_on_non_empty_registry(self, pickle):
         mini_me2 = None
         self.is_empty()
 
-    def dummy(self):
-        pass
-
     def test_pickling(self):
         assert hasattr(pickle.loads(pickle.dumps(cbook.CallbackRegistry())),
                        "callbacks")
