Merge pull request #349 from robotools/notification-fix

Fix behavior in removeObserver.
robotools · Mar 11, 2021 · 9abb77a · 9abb77a
2 parents cfe7998 + 5dccf65
commit 9abb77a
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 12 deletions.
diff --git a/Lib/defcon/test/tools/test_notifications.py b/Lib/defcon/test/tools/test_notifications.py
@@ -71,7 +71,7 @@ def test_removeObserver_notification_observable(self):
         center.removeObserver(observer, "A", observable)
         self.assertFalse(center.hasObserver(observer, "A", observable))
 
-    def test_removeObserver_noNotification(self):
+    def test_removeObserver_allNotifications(self):
         center = NotificationCenter()
         observable1 = _TestObservable(center, "Observable1")
         observable2 = _TestObservable(center, "Observable2")
@@ -85,7 +85,7 @@ def test_removeObserver_noNotification(self):
         center.addObserver(observer1, "notificationCallback", "B", observable2)
         center.addObserver(observer2, "notificationCallback", "A", observable2)
         center.addObserver(observer2, "notificationCallback", "B", observable2)
-        center.removeObserver(observer1, None, observable1)
+        center.removeObserver(observer1, "all", observable1)
         self.assertFalse(center.hasObserver(observer1, "A", observable1))
         self.assertFalse(center.hasObserver(observer1, "B", observable1))
         self.assertTrue(center.hasObserver(observer1, "A", observable2))
@@ -108,8 +108,10 @@ def test_removeObserver_no_notification_observable(self):
         observable = _TestObservable(center, "Observable")
         observer = NotificationTestObserver()
         center.addObserver(observer, "notificationCallback", None, observable)
+        center.addObserver(observer, "notificationCallback", "A", observable)
         center.removeObserver(observer, None, observable)
         self.assertFalse(center.hasObserver(observer, None, observable))
+        self.assertTrue(center.hasObserver(observer, "A", observable))
 
     def test_removeObserver_no_notification_no_observable(self):
         center = NotificationCenter()

diff --git a/Lib/defcon/tools/notifications.py b/Lib/defcon/tools/notifications.py
@@ -49,6 +49,7 @@ class NotificationCenter(object):
 
     def __init__(self):
         self._registry = {}
+        self._observerKeyBacktrack = {}
         self._holds = {}
         self._disabled = {}
         self._identifierRegistry = {}
@@ -83,6 +84,7 @@ def addObserver(self, observer, methodName, notification=None, observable=None,
         must accept a single *notification* argument. This will
         be a :class:`Notification` object.
         """
+        assert notification != "all", "'all' is a reserved name and can not be used as a notification."
         if observable is not None:
             observable = weakref.ref(observable)
         observer = weakref.ref(observer)
@@ -94,6 +96,11 @@ def addObserver(self, observer, methodName, notification=None, observable=None,
                 notification=key[0], observable=key[1](), observer=observer(), method1=self._registry[key][observer], method2=methodName
             )
         self._registry[key][observer] = methodName
+        if observer not in self._observerKeyBacktrack:
+            self._observerKeyBacktrack[observer] = {}
+        if observable not in self._observerKeyBacktrack[observer]:
+            self._observerKeyBacktrack[observer][observable] = set()
+        self._observerKeyBacktrack[observer][observable].add(key)
         if identifier is not None:
             if key not in self._identifierRegistry:
                 self._identifierRegistry[key] = OrderedDict()
@@ -119,22 +126,15 @@ def removeObserver(self, observer, notification, observable=None):
 
         * **observer** A registered object.
         * **notification** The notification that the observer was registered
-          to be notified of. If this is None, all notifications for
+          to be notified of. If this is `"all"`, all notifications for
           the *observable* will be removed for *observer*.
         * **observable** The object being observed.
         """
         if observable is not None:
             observable = weakref.ref(observable)
         observer = weakref.ref(observer)
-        if notification is None:
-            keys = []
-            for (otherNotification, otherObservable), observerDict in self._registry.items():
-                if otherObservable != observable:
-                    continue
-                for otherObserver in observerDict.keys():
-                    if otherObserver != observer:
-                        continue
-                    keys.append((otherNotification, observable))
+        if notification == "all":
+            keys = set(self._observerKeyBacktrack[observer][observable])
         else:
             keys = [
                 (notification, observable)
@@ -146,6 +146,14 @@ def removeObserver(self, observer, notification, observable=None):
                 del self._registry[key][observer]
             if not len(self._registry[key]):
                 del self._registry[key]
+            if observer in self._observerKeyBacktrack:
+                if observable in self._observerKeyBacktrack[observer]:
+                    if key in self._observerKeyBacktrack[observer][observable]:
+                        self._observerKeyBacktrack[observer][observable].remove(key)
+                    if not self._observerKeyBacktrack[observer][observable]:
+                        del self._observerKeyBacktrack[observer][observable]
+                if not self._observerKeyBacktrack[observer]:
+                    del self._observerKeyBacktrack[observer]
             if key in self._identifierRegistry:
                 if observer in self._identifierRegistry[key]:
                     del self._identifierRegistry[key][observer]