From 3b140786467cd01ca1f355b89b55a0d27838ca16 Mon Sep 17 00:00:00 2001
From: Josh Wu <joshwu@google.com>
Date: Thu, 21 Sep 2023 02:04:30 +0800
Subject: [PATCH] Overload signatures

---
 bumble/utils.py | 42 ++++++++++++++++++++++++++++++++++++++----
 1 file 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)