Update participant.py

Author: theomonnom

livekit-rtc/livekit/rtc/participant.py

@@ -19,7 +19,7 @@
 import os
 import mimetypes
 import aiofiles
-from typing import List, Union, Callable, Dict, Awaitable, Optional, Mapping, cast
+from typing import List, Union, Callable, Dict, Awaitable, Optional, Mapping, cast, TypeVar
 from abc import abstractmethod, ABC
 
 from ._ffi_client import FfiClient, FfiHandle
@@ -144,6 +144,10 @@ def disconnect_reason(
         return self._info.disconnect_reason
 
 
+RpcHandler = Callable[["RpcInvocationData"], Union[Awaitable[Optional[str]], Optional[str]]]
+F = TypeVar("F", bound=RpcHandler)
+
+
 class LocalParticipant(Participant):
     """Represents the local participant in a room."""
 
@@ -155,9 +159,7 @@ def __init__(
         super().__init__(owned_info)
         self._room_queue = room_queue
         self._track_publications: dict[str, LocalTrackPublication] = {}  # type: ignore
-        self._rpc_handlers: Dict[
-            str, Callable[[RpcInvocationData], Union[Awaitable[str], str]]
-        ] = {}
+        self._rpc_handlers: Dict[str, RpcHandler] = {}
 
     @property
     def track_publications(self) -> Mapping[str, LocalTrackPublication]:
@@ -328,8 +330,8 @@ async def perform_rpc(
     def register_rpc_method(
         self,
         method_name: str,
-        handler: Optional[Callable[[RpcInvocationData], Union[Awaitable[str], str]]] = None,
-    ) -> Union[None, Callable]:
+        handler: Optional[F] = None,
+    ) -> Union[F, Callable[[F], F]]:
         """
         Establishes the participant as a receiver for calls of the specified RPC method.
         Can be used either as a decorator or a regular method.
@@ -366,18 +368,17 @@ async def greet_handler(data: RpcInvocationData) -> str:
             room.local_participant.register_rpc_method('greet', greet_handler)
         """
 
-        def register(handler_func):
+        def register(handler_func: F) -> F:
             self._rpc_handlers[method_name] = handler_func
             req = proto_ffi.FfiRequest()
             req.register_rpc_method.local_participant_handle = self._ffi_handle.handle
             req.register_rpc_method.method = method_name
             FfiClient.instance.request(req)
+            return handler_func
 
         if handler is not None:
-            register(handler)
-            return None
+            return register(handler)
         else:
-            # Called as a decorator
             return register
 
     def unregister_rpc_method(self, method: str) -> None: