fix(rtc): pass once() handlers only the arguments they take, and let off() remove them

livekit-rtc/livekit/rtc/event_emitter.py

@@ -16,6 +16,43 @@
 T_contra = TypeVar("T_contra", contravariant=True)
 
 
+class _OnceHandler:
+    """The handler once() registers: removes itself, then calls the callback."""
+
+    def __init__(self, emitter: "EventEmitter[Any]", event: Any, callback: Callable) -> None:
+        self._emitter = emitter
+        self._event = event
+        self.callback = callback
+
+    def __call__(self, *args: Any, **kwargs: Any) -> None:
+        self._emitter.off(self._event, self)
+        self.callback(*args, **kwargs)
+
+
+def _signature(handler: Callable) -> inspect.Signature:
+    """Signature emit() trims the arguments to: a once() handler's callback's, read at call time."""
+    if isinstance(handler, _OnceHandler):
+        try:
+            return inspect.signature(handler.callback)
+        except (TypeError, ValueError):
+            pass  # no signature (e.g. a builtin): the handler's *args passes everything
+    return inspect.signature(handler)
+
+
+def _is_once_handler_for(handler: Callable, callback: Callable) -> bool:
+    """Whether handler is a once() handler of callback, matched the way a dict key is."""
+    if not isinstance(handler, _OnceHandler):
+        return False
+    target = handler.callback
+    if target is callback:
+        return True
+    try:
+        # compare hashes first, so an unrelated callback's __eq__ is never called
+        return hash(target) == hash(callback) and bool(target == callback)
+    except TypeError:  # unhashable
+        return False
+
+
 class EventEmitter(Generic[T_contra]):
     def __init__(self) -> None:
         """
@@ -51,7 +88,7 @@ def greet(name):
             callables = list(self._events[event])
             for callback in callables:
                 try:
-                    sig = inspect.signature(callback)
+                    sig = _signature(callback)
                     params = sig.parameters.values()
 
                     has_varargs = any(p.kind == p.VAR_POSITIONAL for p in params)
@@ -115,12 +152,9 @@ def greet_once(name):
             ```
         """
         if callback is not None:
-
-            def once_callback(*args: Any, **kwargs: Any) -> None:
-                self.off(event, once_callback)
-                callback(*args, **kwargs)
-
-            return self.on(event, once_callback)
+            # emit() trims the arguments to the callback's signature, as for on(), and
+            # off() removes the handler by the callback it was registered with
+            return self.on(event, _OnceHandler(self, event, callback))
         else:
 
             def decorator(callback: Callable) -> Callable:
@@ -212,4 +246,8 @@ def greet(name):
             ```
         """
         if event in self._events:
-            self._events[event].pop(callback, None)
+            handlers = self._events[event]
+            handlers.pop(callback, None)
+            once_handlers = [h for h in handlers if _is_once_handler_for(h, callback)]
+            for handler in once_handlers:
+                handlers.pop(handler)

tests/rtc/test_emitter.py

@@ -287,3 +287,104 @@ def handler() -> None:
     emitter.on("event", handler)
     emitter.emit("event")
     assert calls == ["x"]
+
+
+def test_once_passes_only_the_arguments_the_callback_takes() -> None:
+    emitter = EventEmitter[str]()
+    calls: list[Any] = []
+
+    @emitter.once("event")
+    def on_event(first: Any) -> None:
+        calls.append(first)
+
+    emitter.emit("event", 1, 2, 3)  # like on(), the extra arguments are dropped
+    emitter.emit("event", 4)
+
+    assert calls == [1]
+
+
+def test_once_runs_a_callback_without_a_signature() -> None:
+    emitter = EventEmitter[str]()
+    results: list[Any] = []
+
+    emitter.once("event", min)  # builtins like min expose no signature
+    emitter.once("event", lambda *args: results.append(args))
+    emitter.emit("event", 2, 1)
+    emitter.emit("event", 3, 4)
+
+    assert results == [(2, 1)]
+    assert not emitter._events["event"]
+
+
+def test_off_removes_a_once_handler() -> None:
+    emitter = EventEmitter[str]()
+    calls: list[str] = []
+
+    def handler() -> None:
+        calls.append("x")
+
+    class Listener:
+        def handle(self) -> None:
+            calls.append("method")
+
+    listener = Listener()
+    emitter.once("event", handler)
+    emitter.once("event", listener.handle)
+    emitter.off("event", handler)
+    emitter.off("event", listener.handle)  # a new bound method object, equal to the first
+    emitter.emit("event")
+
+    assert calls == []
+
+
+def test_once_trims_to_the_callback_signature_at_emit_time() -> None:
+    emitter = EventEmitter[str]()
+    calls: list[tuple[Any, Any]] = []
+
+    def handler(a: Any, b: Any) -> None:
+        calls.append((a, b))
+
+    callback = functools.partial(handler, b="registered")
+    emitter.once("event", callback)
+    callback.keywords.clear()  # a partial's signature can change after registration
+    emitter.emit("event", 1, 2)
+
+    assert calls == [(1, 2)]
+
+
+def test_off_does_not_compare_unrelated_once_callbacks() -> None:
+    emitter = EventEmitter[str]()
+
+    class Strict:
+        def __call__(self) -> None:
+            pass
+
+        def __eq__(self, other: object) -> bool:
+            if not isinstance(other, Strict):
+                raise ValueError("compared with an unrelated callable")
+            return other is self
+
+        __hash__ = object.__hash__
+
+    def plain() -> None:
+        pass
+
+    emitter.once("event", Strict())
+    emitter.on("event", plain)
+    emitter.off("event", plain)
+
+    assert len(emitter._events["event"]) == 1
+
+
+def test_an_on_handler_with_a_once_callback_attribute_keeps_its_signature() -> None:
+    emitter = EventEmitter[str]()
+    calls: list[Any] = []
+
+    def handler(value: Any) -> None:
+        calls.append(value)
+
+    handler._once_callback = lambda: None  # type: ignore[attr-defined]
+    emitter.on("event", handler)
+    emitter.emit("event", 1)
+
+    assert calls == [1]