From c2fffafa4ef2ad8463b758e8554f11cae9472e5d Mon Sep 17 00:00:00 2001 From: Gergely Polonkai Date: Sat, 18 Nov 2023 07:09:19 +0100 Subject: [PATCH] feat: Create the PSHandler type --- ssb/packet_stream.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/ssb/packet_stream.py b/ssb/packet_stream.py index ade07a1..2d96827 100644 --- a/ssb/packet_stream.py +++ b/ssb/packet_stream.py @@ -34,6 +34,7 @@ from secret_handshake.network import SHSDuplexStream import simplejson from typing_extensions import Self +PSHandler = Union["PSRequestHandler", "PSStreamHandler"] PSMessageData = Union[bytes, bool, Dict[str, Any], str] logger = logging.getLogger("packet_stream") @@ -183,10 +184,10 @@ class PacketStream: def __init__(self, connection: SHSDuplexStream): self.connection = connection self.req_counter = 1 - self._event_map: Dict[int, Tuple[float, Union[PSRequestHandler, PSStreamHandler]]] = {} + self._event_map: Dict[int, Tuple[float, PSHandler]] = {} self._connected = False - def register_handler(self, handler: Union[PSRequestHandler, PSStreamHandler]) -> None: + def register_handler(self, handler: PSHandler) -> None: """Register an RPC handler""" self._event_map[handler.req] = (time(), handler) @@ -283,7 +284,7 @@ class PacketStream: stream: bool = False, end_err: bool = False, req: Optional[int] = None, - ) -> Union[PSRequestHandler, PSStreamHandler]: + ) -> PSHandler: """Send data through the packet stream""" update_counter = False @@ -298,7 +299,7 @@ class PacketStream: self._write(msg) if stream: - handler: Union[PSRequestHandler, PSStreamHandler] = PSStreamHandler(self.req_counter) + handler: PSHandler = PSStreamHandler(self.req_counter) else: handler = PSRequestHandler(self.req_counter)