'with' statement for PubSub (#765)

PubSub objects are now context managers.

2 files changed, 15 insertions, 0 deletions

diff --git a/redis/client.py b/redis/client.py
index 93da07c..0486022 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -3339,6 +3339,12 @@ class PubSub(object):
             ]
         self.reset()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.reset()
+
     def __del__(self):
         try:
             # if this object went out of scope prior to shutting down
diff --git a/tests/test_pubsub.py b/tests/test_pubsub.py
index 2644d60..31b60be 100644
--- a/tests/test_pubsub.py
+++ b/tests/test_pubsub.py
@@ -448,6 +448,15 @@ class TestPubSubAutoDecoding(object):
                                                  new_data,
                                                  pattern=self.pattern)
 
+    def test_context_manager(self, r):
+        with r.pubsub() as pubsub:
+            pubsub.subscribe('foo')
+            assert pubsub.connection is not None
+
+        assert pubsub.connection is None
+        assert pubsub.channels == {}
+        assert pubsub.patterns == {}
+
 
 class TestPubSubRedisDown(object):