From a45be79d02c2eb1f7983ecc722fbdcde5e3008b8 Mon Sep 17 00:00:00 2001
From: Stefano Pigozzi <ste.pigozzi@gmail.com>
Date: Sun, 20 Oct 2019 18:28:47 +0200
Subject: [PATCH] 5.1b1

---
 royalherald/broadcast.py |  8 +++++++-
 royalherald/link.py      | 20 ++++++++++++--------
 royalherald/request.py   |  8 +++++++-
 setup.py                 |  2 +-
 4 files changed, 27 insertions(+), 11 deletions(-)

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",