feat: Create the PSHandler type

This commit is contained in:
Gergely Polonkai 2023-11-18 07:09:19 +01:00
parent 55abb6e341
commit c2fffafa4e
No known key found for this signature in database
GPG Key ID: 2D2885533B869ED4

View File

@ -34,6 +34,7 @@ from secret_handshake.network import SHSDuplexStream
import simplejson import simplejson
from typing_extensions import Self from typing_extensions import Self
PSHandler = Union["PSRequestHandler", "PSStreamHandler"]
PSMessageData = Union[bytes, bool, Dict[str, Any], str] PSMessageData = Union[bytes, bool, Dict[str, Any], str]
logger = logging.getLogger("packet_stream") logger = logging.getLogger("packet_stream")
@ -183,10 +184,10 @@ class PacketStream:
def __init__(self, connection: SHSDuplexStream): def __init__(self, connection: SHSDuplexStream):
self.connection = connection self.connection = connection
self.req_counter = 1 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 self._connected = False
def register_handler(self, handler: Union[PSRequestHandler, PSStreamHandler]) -> None: def register_handler(self, handler: PSHandler) -> None:
"""Register an RPC handler""" """Register an RPC handler"""
self._event_map[handler.req] = (time(), handler) self._event_map[handler.req] = (time(), handler)
@ -283,7 +284,7 @@ class PacketStream:
stream: bool = False, stream: bool = False,
end_err: bool = False, end_err: bool = False,
req: Optional[int] = None, req: Optional[int] = None,
) -> Union[PSRequestHandler, PSStreamHandler]: ) -> PSHandler:
"""Send data through the packet stream""" """Send data through the packet stream"""
update_counter = False update_counter = False
@ -298,7 +299,7 @@ class PacketStream:
self._write(msg) self._write(msg)
if stream: if stream:
handler: Union[PSRequestHandler, PSStreamHandler] = PSStreamHandler(self.req_counter) handler: PSHandler = PSStreamHandler(self.req_counter)
else: else:
handler = PSRequestHandler(self.req_counter) handler = PSRequestHandler(self.req_counter)