Fixed memory leaks.

3 files changed, 52 insertions, 2 deletions

diff --git a/CHANGES b/CHANGES
index b4c0563..761c250 100644
--- a/CHANGES
+++ b/CHANGES
@@ -2,6 +2,10 @@
 Blinker Changelog
 =================
 
+Version 1.4.dev
+---------------
+
+- Fixed memory leaks.
 
 Version 1.3
 -----------
diff --git a/blinker/base.py b/blinker/base.py
index 8fa0390..675c2e0 100644
--- a/blinker/base.py
+++ b/blinker/base.py
@@ -329,11 +329,21 @@ class Signal(object):
     def _disconnect(self, receiver_id, sender_id):
         if sender_id == ANY_ID:
             if self._by_receiver.pop(receiver_id, False):
-                for bucket in self._by_sender.values():
+                empty_buckets = []
+                for key, bucket in self._by_sender.items():
                     bucket.discard(receiver_id)
+                    if not bucket:
+                        empty_buckets.append(key)
+                for key in empty_buckets:
+                    del self._by_sender[key]
             self.receivers.pop(receiver_id, None)
         else:
             self._by_sender[sender_id].discard(receiver_id)
+            if not self._by_sender[sender_id]:
+                del self._by_sender[sender_id]
+            self._by_receiver[receiver_id].discard(sender_id)
+            if not self._by_receiver[receiver_id]:
+                del self._by_receiver[receiver_id]
 
     def _cleanup_receiver(self, receiver_ref):
         """Disconnect a receiver from all senders."""
@@ -346,6 +356,8 @@ class Signal(object):
         self._weak_senders.pop(sender_id, None)
         for receiver_id in self._by_sender.pop(sender_id, ()):
             self._by_receiver[receiver_id].discard(sender_id)
+            if not self._by_receiver[receiver_id]:
+                del self._by_receiver[receiver_id]
 
     def _clear_state(self):
         """Throw away all signal state.  Useful for unit tests."""
diff --git a/tests/test_signals.py b/tests/test_signals.py
index e6142fb..e4fc24d 100644
--- a/tests/test_signals.py
+++ b/tests/test_signals.py
@@ -180,6 +180,40 @@ def test_signal_signals_weak_sender():
     assert len(sentinel) == 1
 
 
+def test_memory_leaks():
+    sentinel = Sentinel()
+    sig = blinker.Signal()
+    receiver1 = sentinel.make_receiver('receiver1')
+    receiver2 = sentinel.make_receiver('receiver2')
+
+    class Sender(object):
+        """A weakref-able object."""
+
+    sender = Sender()
+    sig.connect(receiver1, sender=sender)
+    sig.connect(receiver2, sender=sender)
+
+    assert len(sig._by_sender) == 1
+    assert len(sig._by_receiver) == 2
+    receivers_bucket = list(sig._by_sender.values())[0]
+    assert len(receivers_bucket) == 2
+    senders_buckets = list(sig._by_receiver.values())
+    assert len(senders_buckets[0]) == 1
+    assert len(senders_buckets[1]) == 1
+
+    sig.disconnect(receiver1, sender=sender)
+    assert len(sig._by_sender) == 1
+    assert len(sig._by_receiver) == 1
+    receivers_bucket = list(sig._by_sender.values())[0]
+    assert len(receivers_bucket) == 1
+    senders_bucket = list(sig._by_receiver.values())[0]
+    assert len(senders_bucket) == 1
+
+    sig.disconnect(receiver2, sender=sender)
+    assert len(sig._by_sender) == 0
+    assert len(sig._by_receiver) == 0
+
+
 def test_meta_connect_failure():
     def meta_received(sender, **kw):
         raise TypeError('boom')
@@ -194,7 +228,7 @@ def test_meta_connect_failure():
     assert_raises(TypeError, sig.connect, receiver)
     assert not sig.receivers
     assert not sig._by_receiver
-    assert sig._by_sender == {blinker.base.ANY_ID: set()}
+    assert not sig._by_sender
 
     blinker.receiver_connected._clear_state()