diff --git a/livekit-rtc/livekit/rtc/event_emitter.py b/livekit-rtc/livekit/rtc/event_emitter.py index 38f54901..39364ea9 100644 --- a/livekit-rtc/livekit/rtc/event_emitter.py +++ b/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) diff --git a/tests/rtc/test_emitter.py b/tests/rtc/test_emitter.py index 4f67f76a..289875bb 100644 --- a/tests/rtc/test_emitter.py +++ b/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]