From b50a48fed3b8e90f6a60fb6bf5660f61d43886c6 Mon Sep 17 00:00:00 2001 From: Gilles Boccon-Gibod Date: Tue, 19 Sep 2023 16:14:53 -0700 Subject: [PATCH] add pyee group util --- bumble/utils.py | 90 ++++++++++++++++++++++++++++++-- tests/utils_test.py | 122 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 209 insertions(+), 3 deletions(-) create mode 100644 tests/utils_test.py diff --git a/bumble/utils.py b/bumble/utils.py index 8a55684..cddd4e1 100644 --- a/bumble/utils.py +++ b/bumble/utils.py @@ -19,9 +19,19 @@ import asyncio import logging import traceback import collections +from contextlib import ExitStack import sys -from typing import Awaitable, Set, TypeVar -from functools import wraps +from typing import ( + Awaitable, + Callable, + ContextManager, + Mapping, + Optional, + Set, + Tuple, + TypeVar, +) +import functools from pyee import EventEmitter from .colors import color @@ -167,7 +177,7 @@ class AsyncRunner: """ def decorator(func): - @wraps(func) + @functools.wraps(func) def wrapper(*args, **kwargs): coroutine = func(*args, **kwargs) if queue is None: @@ -302,3 +312,77 @@ class FlowControlAsyncPipe: self.resume_source() self.check_pump() + + +# ----------------------------------------------------------------------------- +def event_emitter_once_for_group( + emitter: EventEmitter, + handlers: Mapping[str, Callable], + context: Optional[ContextManager] = None, +) -> None: + """ + Register event listeners as a group, with optional context manager. + + For each entry in the map, an event listener is registered with the emitter. + When any of the event names in the handlers map is emitted by the emitter, + the corresponding handler is invoked, then all of the listeners are removed from + the emitter. + If a context manager is specified, it will be entered prior to registering the + listeners, and exited when any of the events is emitted. + + Args: + emitter: + The event emitter with which to register the event listeners. + handlers: + A map that associates an event name with an event handler. + context: + A context manager to manager resources, or None if not needed. + """ + event_emitters_once_for_group( + {(emitter, event_name): handler for event_name, handler in handlers.items()}, + context, + ) + + +# ----------------------------------------------------------------------------- +def event_emitters_once_for_group( + handlers: Mapping[Tuple[EventEmitter, str], Callable], + context: Optional[ContextManager] = None, +) -> None: + """ + Register event listeners as a group, with optional context manager. + + Similar to event_emitter_once_for_group, but instead of registering the listeners + with a single emitter, each event may by associate with a different emitter. + + Args: + handlers: + A map that associates an (emitter, event_name) pair with an event handler. + context: + A context manager to manager resources, or None if not needed. + """ + # Setup an exit stack to enter and exit the context, if any. + if context is not None: + exit_stack = ExitStack() + exit_stack.enter_context(context) + else: + exit_stack = None + + def on_event(handler, *args, **kwargs) -> None: + # Invoke the handler. + handler(*args, **kwargs) + + # Release the context, if any. + if exit_stack is not None: + exit_stack.close() + + # Remove all listeners. + for (emitter, event_name), listener in listeners.items(): + emitter.remove_listener(event_name, listener) + + listeners = { + (emitter, event_name): emitter.on( + event_name, functools.partial(on_event, handler) + ) + for (emitter, event_name), handler in handlers.items() + } diff --git a/tests/utils_test.py b/tests/utils_test.py new file mode 100644 index 0000000..adc35bc --- /dev/null +++ b/tests/utils_test.py @@ -0,0 +1,122 @@ +# Copyright 2021-2022 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ----------------------------------------------------------------------------- +# Imports +# ----------------------------------------------------------------------------- +import contextlib +from pyee import EventEmitter + +from bumble.utils import event_emitter_once_for_group, event_emitters_once_for_group + + +# ----------------------------------------------------------------------------- +def test_event_emitter_once_for_group(): + results = {'event1': None, 'event2': None, 'released': 0} + + def handler1(x): + results['event1'] = x + + def handler2(y): + results['event2'] = y + + emitter = EventEmitter() + + event_emitter_once_for_group( + emitter, + { + 'event1': handler1, + 'event2': handler2, + }, + ) + + emitter.emit('event1', 'hello') + + assert results['event1'] == 'hello' + assert results['event2'] is None + + results['event1'] = None + + emitter.emit('event1', 'hello') + emitter.emit('event2', 1234) + + assert results['event1'] is None + assert results['event2'] is None + + @contextlib.contextmanager + def managed(): + try: + yield 1234 + finally: + results['released'] += 1 + + event_emitter_once_for_group( + emitter, + { + 'event1': handler1, + 'event2': handler2, + }, + managed(), + ) + + assert results['released'] == 0 + + emitter.emit('event2', 7756) + + assert results['event1'] is None + assert results['event2'] == 7756 + assert results['released'] == 1 + + +# ----------------------------------------------------------------------------- +def test_event_emitters_once_for_group(): + results = {'event1': None, 'event2': None, 'released': 0} + + def handler1(x): + results['event1'] = x + + def handler2(y): + results['event2'] = y + + emitter1 = EventEmitter() + emitter2 = EventEmitter() + + event_emitters_once_for_group( + { + (emitter1, 'event1'): handler1, + (emitter2, 'event2'): handler2, + }, + ) + + emitter1.emit('event1', 'hello') + emitter2.emit('event1', 'foobar') + + assert results['event1'] == 'hello' + assert results['event2'] is None + + results['event1'] = None + + emitter1.emit('event1', 'hello') + emitter1.emit('event2', 1234) + emitter2.emit('event1', 'hello') + emitter2.emit('event2', 1234) + + assert results['event1'] is None + assert results['event2'] is None + + +# ----------------------------------------------------------------------------- +if __name__ == '__main__': + test_event_emitter_once_for_group() + test_event_emitters_once_for_group()