ENH: Add the ability to block callback signals

greglucas · greglucas · commit 1cf6384c4c83 · 2021-08-11T08:11:32.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/users/next_whats_new/callback_blocking.rst b/doc/users/next_whats_new/callback_blocking.rst
@@ -0,0 +1,23 @@
+``CallbackRegistry`` objects gain a method to temporarily block signals
+-----------------------------------------------------------------------
+
+The context manager `~matplotlib.cbook.CallbackRegistry.blocked` 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.
+
+.. code-block::
+
+    from matplotlib.cbook import CallbackRegistry
+    callbacks = CallbackRegistry()
+    def test_func():
+        print("Inside test function")
+    callbacks.connect("test", test_func)
+
+    # Pause the callbacks
+    with callbacks.blocked():
+        callbacks.process("test")  # Nothing printed
+    
+    # Pause a single callback if there are multiple connected signals
+    with callbacks.blocked(signal="test"):
+        callbacks.process("test")  # Nothing printed
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)
@@ -140,9 +141,15 @@ class CallbackRegistry:
         >>> callbacks.process('eat', 456)
         eat 456
         >>> callbacks.process('be merry', 456) # nothing will be called
+
         >>> callbacks.disconnect(id_eat)
         >>> callbacks.process('eat', 456)      # nothing will be called
 
+        >>> with callbacks.blocked(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 +287,31 @@ def process(self, s, *args, **kwargs):
                     else:
                         raise
 
+    @contextlib.contextmanager
+    def blocked(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, optional
+            The callback signal to block. The default is to block 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.blocked():
+        cb.process("test1")
+        cb.process("test2")
+
+    # block individual callbacks to make sure the other is still processed
+    with cb.blocked(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']
