Added ``Signal.temporarily_connected_to`` context manager

5 files changed, 109 insertions, 0 deletions

diff --git a/CHANGES b/CHANGES
index 5c6aa82..55b0df1 100644
--- a/CHANGES
+++ b/CHANGES
@@ -2,6 +2,14 @@
 Blinker Changelog
 =================
 
+Version 0.9
+-----------
+
+Not yet released
+
+- Added ``Signal.temporarily_connected_to`` context manager
+
+
 Version 0.8
 -----------
 
diff --git a/blinker/_utilities.py b/blinker/_utilities.py
index 9890194..4cb7cbd 100644
--- a/blinker/_utilities.py
+++ b/blinker/_utilities.py
@@ -57,6 +57,16 @@ except:
                                             dict.__repr__(self))
 
 
+try:
+    from contextlib import contextmanager
+except ImportError:
+    def contextmanager(fn):
+        def oops(*args, **kw):
+            raise RuntimeError("Python 2.5 or above is required to use "
+                               "context managers.")
+        oops.__name__ = fn.__name__
+        return oops
+
 class _symbol(object):
 
     def __init__(self, name):
diff --git a/blinker/base.py b/blinker/base.py
index 14ee6d9..56b1c54 100644
--- a/blinker/base.py
+++ b/blinker/base.py
@@ -12,6 +12,7 @@ from weakref import WeakValueDictionary
 
 from blinker._utilities import (
     WeakTypes,
+    contextmanager,
     defaultdict,
     hashable_identity,
     reference,
@@ -94,6 +95,32 @@ class Signal(object):
                 raise
         return receiver
 
+    @contextmanager
+    def temporarily_connected_to(self, receiver, sender=ANY):
+        """Execute a block with the signal connected *receiver*.
+
+        This is a context manager for use in the ``with`` statement, and can
+        be useful in unit tests.  *receiver* is connected to the signal for
+        the duration of the ``with`` block, and will be disconnected
+        automatically when exiting the block::
+
+          ready = Signal()
+          receiver = lambda sender: pass
+
+          with ready.temporarily_connected_to(receiver):
+             # do stuff
+             ready.send(123)
+
+        """
+        self.connect(receiver, sender=sender, weak=False)
+        try:
+            yield None
+        except:
+            self.disconnect(receiver)
+            raise
+        else:
+            self.disconnect(receiver)
+
     def send(self, *sender, **kwargs):
         """Emit this signal on behalf of *sender*, passing on \*\*kwargs.
 
diff --git a/tests/test_context.py b/tests/test_context.py
new file mode 100644
index 0000000..b5de372
--- /dev/null
+++ b/tests/test_context.py
@@ -0,0 +1,56 @@
+from __future__ import with_statement
+
+from blinker import Signal
+
+
+def test_temp_connection():
+    sig = Signal()
+
+    canary = []
+    receiver = lambda sender: canary.append(sender)
+
+    sig.send(1)
+    with sig.temporarily_connected_to(receiver):
+        sig.send(2)
+    sig.send(3)
+
+    assert canary == [2]
+    assert not sig.receivers
+
+
+def test_temp_connection_for_sender():
+    sig = Signal()
+
+    canary = []
+    receiver = lambda sender: canary.append(sender)
+
+    with sig.temporarily_connected_to(receiver, sender=2):
+        sig.send(1)
+        sig.send(2)
+
+    assert canary == [2]
+    assert not sig.receivers
+
+
+def test_temp_connection_failure():
+    sig = Signal()
+
+    canary = []
+    receiver = lambda sender: canary.append(sender)
+
+    class Failure(Exception):
+        pass
+
+    try:
+        sig.send(1)
+        with sig.temporarily_connected_to(receiver):
+            sig.send(2)
+            raise Failure
+        sig.send(3)
+    except Failure:
+        pass
+    else:
+        raise AssertionError("Context manager did not propagate.")
+
+    assert canary == [2]
+    assert not sig.receivers
diff --git a/tests/test_signals.py b/tests/test_signals.py
index 9715293..5328872 100644
--- a/tests/test_signals.py
+++ b/tests/test_signals.py
@@ -1,3 +1,4 @@
+import sys
 import blinker
 from nose.tools import assert_raises
 
@@ -233,3 +234,10 @@ def test_named_blinker():
 def values_are_empty_sets_(dictionary):
     for val in dictionary.values():
         assert val == set()
+
+if sys.version_info < (2, 5):
+    def test_context_manager_warning():
+        sig = blinker.Signal()
+        receiver = lambda sender: None
+
+        assert_raises(RuntimeError, sig.temporarily_connected_to, receiver)