ENH: Add the ability to block callback signals

greglucas · greglucas · commit a9680e7332eb · 2021-08-10T16:40:10.000-06:00
CallbackRegistry objects now have a context manager .block()
method that can be used to temporarily block callback signals
from being processed.
diff --git a/doc/api/next_api_changes/behavior/20816-GL.rst b/doc/api/next_api_changes/behavior/20816-GL.rst
@@ -0,0 +1,7 @@
+``CallbackRegistry`` objects gain a method to temporarily block signals
+-----------------------------------------------------------------------
+
+The method `~matplotlib.cbook.CallbackRegistry.block` can be used to block
+callback signals from being processed by the ``CallbackRegistry``.
+The optional keyword, *signal*, can be used to block a specific signal
+from being processed and let all other signals pass.
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -122,7 +122,8 @@ def _weak_or_strong_ref(func, callback):
 
 class CallbackRegistry:
     """
-    Handle registering and disconnecting for a set of signals and callbacks:
+    Handle registering, processing, blocking, and disconnecting
+    for a set of signals and callbacks:
 
         >>> def oneat(x):
         ...    print('eat', x)
@@ -143,6 +144,11 @@ class CallbackRegistry:
         >>> callbacks.disconnect(id_eat)
         >>> callbacks.process('eat', 456)      # nothing will be called
 
+        >>> with callbacks.block(signal='drink'):
+        ...     callbacks.process('drink', 123) # nothing will be called
+        >>> callbacks.process('drink', 123)
+        drink 123
+
     In practice, one should always disconnect all callbacks when they are
     no longer needed to avoid dangling references (and thus memory leaks).
     However, real code in Matplotlib rarely does so, and due to its design,
@@ -280,6 +286,31 @@ def process(self, s, *args, **kwargs):
                     else:
                         raise
 
+    @contextlib.contextmanager
+    def block(self, *, signal=None):
+        """
+        Block callback signals from being processed.
+
+        A context manager to temporarily block/disable callback signals
+        from being processed by the registered listeners.
+
+        Parameters
+        ----------
+        signal : str, default: None
+            The callback signal to block, default blocks all signals
+        """
+        orig = self.callbacks
+        try:
+            if signal is None:
+                # Empty out the callbacks
+                self.callbacks = {}
+            else:
+                # Only remove the specific signal
+                self.callbacks = {k: orig[k] for k in orig if k != signal}
+            yield
+        finally:
+            self.callbacks = orig
+
 
 class silent_list(list):
     """
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -361,6 +361,39 @@ def test_callbackregistry_custom_exception_handler(monkeypatch, cb, excp):
         cb.process('foo')
 
 
+def test_callbackregistry_blocking():
+    # Needs an exception handler for interactive testing environments
+    # that would only print this out instead of raising the exception
+    def raise_handler(excp):
+        raise excp
+    cb = cbook.CallbackRegistry(exception_handler=raise_handler)
+    def test_func1():
+        raise ValueError("1 should be blocked")
+    def test_func2():
+        raise ValueError("2 should be blocked")
+    cb.connect("test1", test_func1)
+    cb.connect("test2", test_func2)
+
+    # block all of the callbacks to make sure they aren't processed
+    with cb.block():
+        cb.process("test1")
+        cb.process("test2")
+
+    # block individual callbacks to make sure the other is still processed
+    with cb.block(signal="test1"):
+        # Blocked
+        cb.process("test1")
+        # Should raise
+        with pytest.raises(ValueError, match="2 should be blocked"):
+            cb.process("test2")
+
+    # Make sure the original callback functions are there after blocking
+    with pytest.raises(ValueError, match="1 should be blocked"):
+        cb.process("test1")
+    with pytest.raises(ValueError, match="2 should be blocked"):
+        cb.process("test2")
+
+
 def test_sanitize_sequence():
     d = {'a': 1, 'b': 2, 'c': 3}
     k = ['a', 'b', 'c']
