diff --git a/royalherald/broadcast.py b/royalherald/broadcast.py index 7e6ae5cd..848aff9b 100644 --- a/royalherald/broadcast.py +++ b/royalherald/broadcast.py @@ -1,6 +1,12 @@ +import typing + + class Broadcast: - def __init__(self, handler: str, data: dict): + def __init__(self, handler: str, data: dict, msg_type: typing.Optional[str] = None): super().__init__() + if msg_type is not None: + assert msg_type == self.__class__.__name__ + self.msg_type = self.__class__.__name__ self.handler: str = handler self.data: dict = data diff --git a/royalherald/link.py b/royalherald/link.py index b060137d..e0f5d586 100644 --- a/royalherald/link.py +++ b/royalherald/link.py @@ -59,8 +59,8 @@ class Link: self.nid: str = str(uuid.uuid4()) self.secret: str = secret self.websocket: typing.Optional[websockets.WebSocketClientProtocol] = None - # Not sure on the type annotation here - self.request_handler: typing.Callable[[Request], typing.Awaitable[Response]] = request_handler + self.request_handler: typing.Callable[[typing.Union[Request, Broadcast]], + typing.Awaitable[Response]] = request_handler self._pending_requests: typing.Dict[str, PendingRequest] = {} if loop is None: self._loop = asyncio.get_event_loop() @@ -170,9 +170,13 @@ class Link: request.set(package.data) continue # Package is a request - assert isinstance(package, Package) - log.debug(f"Received request {package.source_conv_id}: {package}") - response: Response = await self.request_handler(Request.from_dict(package.data)) - response_package: Package = package.reply(response.to_dict()) - await self.send(response_package) - log.debug(f"Replied to request {response_package.source_conv_id}: {response_package}") + elif package.data["msg_type"] == "Request": + log.debug(f"Received request {package.source_conv_id}: {package}") + response: Response = await self.request_handler(Request.from_dict(package.data)) + response_package: Package = package.reply(response.to_dict()) + await self.send(response_package) + log.debug(f"Replied to request {response_package.source_conv_id}: {response_package}") + # Package is a broadcast + elif package.data["msg_type"] == "Broadcast": + log.debug(f"Received broadcast {package.source_conv_id}: {package}") + await self.request_handler(Broadcast.from_dict(package.data)) diff --git a/royalherald/request.py b/royalherald/request.py index 97d8e84b..9e760e67 100644 --- a/royalherald/request.py +++ b/royalherald/request.py @@ -1,10 +1,16 @@ +import typing + + class Request: """A request sent from a :py:class:`Link` to another. It contains the name of the requested handler, in addition to the data.""" - def __init__(self, handler: str, data: dict): + def __init__(self, handler: str, data: dict, msg_type: typing.Optional[str] = None): super().__init__() + if msg_type is not None: + assert msg_type == self.__class__.__name__ + self.msg_type = self.__class__.__name__ self.handler: str = handler self.data: dict = data diff --git a/setup.py b/setup.py index c77692bc..3306892d 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ with open("README.md", "r") as f: setuptools.setup( name="royalherald", - version="5.0b9", + version="5.1b1", author="Stefano Pigozzi", author_email="ste.pigozzi@gmail.com", description="A websocket communication protocol",