Allow different callbacks for different message templates.

diff --git a/controlpi/messagebus.py b/controlpi/messagebus.py
index 6aa69bb267cab8a29aa2c0fb7af28ec63665894b..54579fedde34739fe2453573ea69b87635dafe09 100644 (file)
--- a/controlpi/messagebus.py
+++ b/controlpi/messagebus.py
@@ -30,11 +30,11 @@ means that it wants to send arbitrary or receive all messages, respectively:
 >>> async def setup(bus):
 ...     bus.register('Logger', 'Test Plugin',
 ...                  [],
-...                  [{}],
+...                  [MessageTemplate({})],
 ...                  callback_for_receiver('Logger'))
 ...     bus.register('Client 1', 'Test Plugin',
-...                  [{'k1': {'type': 'string'}}],
-...                  [{'target': {'const': 'Client 1'}}],
+...                  [MessageTemplate({'k1': {'type': 'string'}})],
+...                  [MessageTemplate({'target': {'const': 'Client 1'}})],
 ...                  callback_for_receiver('Client 1'))
 
 While most clients should always use their own name for sending, this is not
@@ -87,7 +87,7 @@ import fastjsonschema  # type: ignore
 import sys
 
 from typing import (Union, Dict, List, Any, Callable, Coroutine,
-                    Optional, Iterable)
+                    Optional, Iterable, Tuple)
 MessageValue = Union[None, str, int, float, bool, Dict[str, Any], List[Any]]
 # Should really be:
 # MessageValue = Union[None, str, int, float, bool,
@@ -1016,15 +1016,15 @@ class MessageBus:
     ...     print("Setting up.")
     ...     bus.register('Logger', 'Test Plugin',
     ...                  [],
-    ...                  [{}],
+    ...                  [MessageTemplate({})],
     ...                  callback_for_receiver('Logger'))
     ...     bus.register('Client 1', 'Test Plugin',
-    ...                  [{'k1': {'type': 'string'}}],
-    ...                  [{'target': {'const': 'Client 1'}}],
+    ...                  [MessageTemplate({'k1': {'type': 'string'}})],
+    ...                  [MessageTemplate({'target': {'const': 'Client 1'}})],
     ...                  callback_for_receiver('Client 1'))
     ...     bus.register('Client 2', 'Test Plugin',
-    ...                  [{}],
-    ...                  [{'target': {'const': 'Client 2'}}],
+    ...                  [MessageTemplate({})],
+    ...                  [MessageTemplate({'target': {'const': 'Client 2'}})],
     ...                  callback_for_receiver('Client 2'))
 
     The bus itself is addressed by the empty string. It sends messages for
@@ -1110,8 +1110,10 @@ class MessageBus:
 
     def register(self, client: str, plugin: str,
                  sends: Iterable[MessageTemplate],
-                 receives: Iterable[MessageTemplate],
-                 callback: MessageCallback) -> None:
+                 receives: Union[Iterable[MessageTemplate],
+                                 Iterable[Tuple[Iterable[MessageTemplate],
+                                                MessageCallback]]],
+                 callback: Optional[MessageCallback] = None) -> None:
         """Register a client at the message bus.
 
         >>> async def callback(message):
@@ -1120,17 +1122,19 @@ class MessageBus:
         ...     bus = MessageBus()
         ...     bus.register('Logger', 'Test Plugin',
         ...                  [],    # send nothing
-        ...                  [{}],  # receive everything
+        ...                  [MessageTemplate({})],  # receive everything
         ...                  callback)
         ...     bus.register('Client 1', 'Test Plugin',
-        ...                  [{'k1': {'type': 'string'}}],
+        ...                  [MessageTemplate({'k1': {'type': 'string'}})],
         ...                      # send with key 'k1' and string value
-        ...                  [{'target': {'const': 'Client 1'}}],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 1'}})],
         ...                      # receive for this client
         ...                  callback)
         ...     bus.register('Client 2', 'Test Plugin',
-        ...                  [{}],  # send arbitrary
-        ...                  [{'target': {'const': 'Client 2'}}],
+        ...                  [MessageTemplate({})],  # send arbitrary
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 2'}})],
         ...                      # receive for this client
         ...                  callback)
         >>> asyncio.run(main())
@@ -1148,8 +1152,13 @@ class MessageBus:
         for template in sends:
             self._send_reg.insert(template, client)
         event['sends'] = self._send_reg.get_templates(client)
-        for template in receives:
-            self._recv_reg.insert(template, client, callback)
+        if callback:
+            for template in receives:
+                self._recv_reg.insert(template, client, callback)
+        else:
+            for (templates, callback) in receives:
+                for template in templates:
+                    self._recv_reg.insert(template, client, callback)
         event['receives'] = self._recv_reg.get_templates(client)
         self._queue.put_nowait(event)
 
@@ -1161,8 +1170,9 @@ class MessageBus:
         >>> async def main():
         ...     bus = MessageBus()
         ...     bus.register('Client 1', 'Test Plugin',
-        ...                  [{'k1': {'type': 'string'}}],
-        ...                  [{'target': {'const': 'Client 1'}}],
+        ...                  [MessageTemplate({'k1': {'type': 'string'}})],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 1'}})],
         ...                  callback)
         ...     bus.unregister('Client 1')
         >>> asyncio.run(main())
@@ -1231,12 +1241,14 @@ class MessageBus:
         >>> async def main():
         ...     bus = MessageBus()
         ...     bus.register('Client 1', 'Test Plugin',
-        ...                  [{'k1': {'type': 'string'}}],
-        ...                  [{'target': {'const': 'Client 1'}}],
+        ...                  [MessageTemplate({'k1': {'type': 'string'}})],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 1'}})],
         ...                  callback)
         ...     bus.register('Client 2', 'Test Plugin',
-        ...                  [{}],
-        ...                  [{'target': {'const': 'Client 2'}}],
+        ...                  [MessageTemplate({})],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 2'}})],
         ...                  callback)
         ...     bus_task = asyncio.create_task(bus.run())
         ...     await bus.send({'sender': 'Client 1', 'target': 'Client 2',
@@ -1271,12 +1283,14 @@ class MessageBus:
         >>> async def main():
         ...     bus = MessageBus()
         ...     bus.register('Client 1', 'Test Plugin',
-        ...                  [{'k1': {'type': 'string'}}],
-        ...                  [{'target': {'const': 'Client 1'}}],
+        ...                  [MessageTemplate({'k1': {'type': 'string'}})],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 1'}})],
         ...                  callback)
         ...     bus.register('Client 2', 'Test Plugin',
-        ...                  [{}],
-        ...                  [{'target': {'const': 'Client 2'}}],
+        ...                  [MessageTemplate({})],
+        ...                  [MessageTemplate({'target':
+        ...                                    {'const': 'Client 2'}})],
         ...                  callback)
         ...     bus_task = asyncio.create_task(bus.run())
         ...     bus.send_nowait({'sender': 'Client 1', 'target': 'Client 2',