Merge pull request #18 from Zsailer/remove-listener

Add method to remove modifier

jupyter_events/logger.py

@@ -178,7 +178,7 @@ def add_modifier(
             A callable function/method that executes when the named event occurs.
             This method enforces a string signature for modifiers:
 
-                (schema_id: str, version: int, data: dict) -> dict:
+                (schema_id: str, data: dict) -> dict:
         """
         # Ensure that this is a callable function/method
         if not callable(modifier):
@@ -215,31 +215,31 @@ def remove_modifier(
         self, *, schema_id: str = None, modifier: Callable[[str, dict], dict]
     ) -> None:
         """Remove a modifier from an event or all events.
+
         Parameters
         ----------
         schema_id: str
             If given, remove this modifier only for a specific event type.
-        modifier: Callable[[str, int, dict], dict]
+        modifier: Callable[[str, dict], dict]
+
             The modifier to remove.
         """
         # If schema_id is given remove the modifier from this schema.
         if schema_id:
-            self._modifiers[schema_id].remove(modifier)
+            self._modifiers[schema_id].discard(modifier)
         # If no schema_id is given, remove the modifier from all events.
         else:
-            for modifier_list in self._modifiers.values():
+            for schema_id in self.schemas.schema_ids:
                 # Remove the modifier if it is found in the list.
-                try:
-                    modifier_list.remove(modifier)
-                except ValueError:
-                    pass
+                self._modifiers[schema_id].discard(modifier)
+                self._modifiers[schema_id].discard(modifier)
 
     def add_listener(
         self,
         *,
         modified: bool = True,
         schema_id: Union[str, None] = None,
-        listener: Callable[[str, int, dict], None],
+        listener: Callable[["EventLogger", str, dict], None],
     ):
         """Add a listener (callable) to a registered event.
 
@@ -292,6 +292,33 @@ async def listener_signature(
                 "and the return value."
             )
 
+    def remove_listener(
+        self,
+        *,
+        schema_id: str = None,
+        listener: Callable[["EventLogger", str, dict], None],
+    ) -> None:
+        """Remove a listener from an event or all events.
+
+        Parameters
+        ----------
+        schema_id: str
+            If given, remove this modifier only for a specific event type.
+
+        listener: Callable[[EventLogger, str, dict], dict]
+            The modifier to remove.
+        """
+        # If schema_id is given remove the listener from this schema.
+        if schema_id:
+            self._modified_listeners[schema_id].discard(listener)
+            self._unmodified_listeners[schema_id].discard(listener)
+        # If no schema_id is given, remove the listener from all events.
+        else:
+            for schema_id in self.schemas.schema_ids:
+                # Remove the listener if it is found in the list.
+                self._modified_listeners[schema_id].discard(listener)
+                self._unmodified_listeners[schema_id].discard(listener)
+
     def emit(self, *, schema_id: str, data: dict, timestamp_override=None):
         """
         Record given event with schema has occurred.

jupyter_events/schema_registry.py

@@ -28,6 +28,10 @@ def _add(self, schema_obj: EventSchema):
             )
         self._schemas[schema_obj.id] = schema_obj
 
+    @property
+    def schema_ids(self):
+        return self._schemas.keys()
+
     def register(self, schema: Union[dict, str, EventSchema]) -> EventSchema:
         """Add a valid schema to the registry.
 

tests/test_listeners.py

@@ -40,6 +40,28 @@ async def my_listener(logger: EventLogger, schema_id: str, data: dict) -> None:
     assert len(event_logger._active_listeners) == 0
 
 
+async def test_remove_listener_function(event_logger, schema):
+    global listener_was_called
+    listener_was_called = False
+
+    async def my_listener(logger: EventLogger, schema_id: str, data: dict) -> None:
+        global listener_was_called
+        listener_was_called = True  # type: ignore
+
+    # Add the modifier
+    event_logger.add_listener(schema_id=schema.id, listener=my_listener)
+    event_logger.emit(schema_id=schema.id, data={"prop": "hello, world"})
+    await event_logger.gather_listeners()
+    assert listener_was_called
+
+    # Check that the active listeners are cleaned up.
+    assert len(event_logger._active_listeners) == 0
+
+    event_logger.remove_listener(listener=my_listener)
+    assert len(event_logger._modified_listeners[schema.id]) == 0
+    assert len(event_logger._unmodified_listeners[schema.id]) == 0
+
+
 async def test_bad_listener_function_signature(event_logger, schema):
     async def listener_with_extra_args(
         logger: EventLogger, schema_id: str, data: dict, unknown_arg: dict