diff options
author | Josh Wu <joshwu@google.com> | 2023-09-21 02:04:30 +0800 |
---|---|---|
committer | Lucas Abel <22837557+uael@users.noreply.github.com> | 2023-09-20 23:03:23 +0200 |
commit | 3b140786467cd01ca1f355b89b55a0d27838ca16 (patch) | |
tree | 5817cffa10a0c643f428fca0b7a980c5094eeae6 | |
parent | eb9d5632bc7ea1f33152b84153cfb8eda2375b16 (diff) | |
download | bumble-3b140786467cd01ca1f355b89b55a0d27838ca16.tar.gz |
Overload signatures
-rw-r--r-- | bumble/utils.py | 42 |
1 files changed, 38 insertions, 4 deletions
diff --git a/bumble/utils.py b/bumble/utils.py index b9f2e59..da5836d 100644 --- a/bumble/utils.py +++ b/bumble/utils.py @@ -21,7 +21,18 @@ import logging import traceback import collections import sys -from typing import Awaitable, Set, TypeVar, List, Tuple, Callable, Any, Optional, Union +from typing import ( + Awaitable, + Set, + TypeVar, + List, + Tuple, + Callable, + Any, + Optional, + Union, + overload, +) from functools import wraps from pyee import EventEmitter @@ -65,13 +76,15 @@ def composite_listener(cls): return cls +# ----------------------------------------------------------------------------- _Handler = TypeVar('_Handler', bound=Callable) -# ----------------------------------------------------------------------------- + class EventWatcher: '''A wrapper class to control the lifecycle of event handlers better. Usage: + ``` watcher = EventWatcher() def on_foo(): @@ -84,13 +97,16 @@ class EventWatcher: # Close all event handlers watching through this watcher watcher.close() + ``` As context: + ``` with contextlib.closing(EventWatcher()) as context: @context.on(emitter, 'foo') def on_foo(): ... # on_foo() has been removed here! + ``` ''' handlers: List[Tuple[EventEmitter, str, Callable[..., Any]]] @@ -98,6 +114,14 @@ class EventWatcher: def __init__(self) -> None: self.handlers = [] + @overload + def on(self, emitter: EventEmitter, event: str) -> Callable[[_Handler], _Handler]: + ... + + @overload + def on(self, emitter: EventEmitter, event: str, handler: _Handler) -> _Handler: + ... + def on( self, emitter: EventEmitter, event: str, handler: Optional[_Handler] = None ) -> Union[_Handler, Callable[[_Handler], _Handler]]: @@ -109,12 +133,21 @@ class EventWatcher: handler: (Optional) Event handler. When nothing passed, this method works as a decorator. ''' - def wrapper(f: _Handler): + def wrapper(f: _Handler) -> _Handler: self.handlers.append((emitter, event, f)) emitter.on(event, f) + return f return wrapper if handler is None else wrapper(handler) + @overload + def once(self, emitter: EventEmitter, event: str) -> Callable[[_Handler], _Handler]: + ... + + @overload + def once(self, emitter: EventEmitter, event: str, handler: _Handler) -> _Handler: + ... + def once( self, emitter: EventEmitter, event: str, handler: Optional[_Handler] = None ) -> Union[_Handler, Callable[[_Handler], _Handler]]: @@ -126,9 +159,10 @@ class EventWatcher: handler: (Optional) Event handler. When nothing passed, this method works as a decorator. ''' - def wrapper(f: _Handler): + def wrapper(f: _Handler) -> _Handler: self.handlers.append((emitter, event, f)) emitter.once(event, f) + return f return wrapper if handler is None else wrapper(handler) |