Add topic handling functions to Agent

osbrain/address.py

@@ -389,15 +389,15 @@ class AgentChannel():
     receiver : str
         Second AgentAddress.
     """
-    def __init__(self, kind, receiver, sender):
+    def __init__(self, kind, receiver, sender, uuid=None):
         self.kind = AgentChannelKind(kind)
         self.receiver = receiver
         self.sender = sender
         self.transport = \
             receiver.transport if receiver else sender.transport
         self.serializer = \
             receiver.serializer if receiver else sender.serializer
-        self.uuid = unique_identifier()
+        self.uuid = uuid or unique_identifier()
         # Set up pairs
         if sender:
             self.sender.channel = self
@@ -435,4 +435,5 @@ def twin(self):
         kind = self.kind.twin()
         sender = self.receiver.twin() if self.receiver is not None else None
         receiver = self.sender.twin() if self.sender is not None else None
-        return self.__class__(kind=kind, receiver=receiver, sender=sender)
+        return self.__class__(kind=kind, receiver=receiver, sender=sender,
+                              uuid=self.uuid)

osbrain/agent.py

@@ -898,6 +898,44 @@ def _connect_and_register(self, client_address, alias=None, handler=None,
         self.register(socket, register_as, alias, handler)
         return client_address
 
+    def unsubscribe_socket_from_topic(self, socket, topic: bytes):
+        '''
+        Unsubscribe a socket from a given topic.
+
+        Parameters
+        ----------
+        socket
+            Identifier of the socket. Must be a valid entry of `self.socket`
+        topic
+            topic which we want to unsubscribe from
+        '''
+        if isinstance(self.address[socket], AgentAddress):
+            self.socket[socket].setsockopt(zmq.UNSUBSCRIBE, topic)
+        if isinstance(self.address[socket], AgentChannel):
+            channel = self.address[socket]
+            socket = channel.receiver
+            treated_topic = channel.uuid + topic
+            self.socket[socket].setsockopt(zmq.UNSUBSCRIBE, treated_topic)
+
+    def subscribe_socket_to_topic(self, socket, topic: bytes):
+        '''
+        Subscribe a socket to a given topic.
+
+        Parameters
+        ----------
+        socket
+            Identifier of the socket. Must be a valid entry of `self.socket`
+        topic
+            topic which we want to subscribe to
+        '''
+        if isinstance(self.address[socket], AgentAddress):
+            self.socket[socket].setsockopt(zmq.SUBSCRIBE, topic)
+        if isinstance(self.address[socket], AgentChannel):
+            channel = self.address[socket]
+            socket = channel.receiver
+            treated_topic = channel.uuid + topic
+            self.socket[socket].setsockopt(zmq.SUBSCRIBE, treated_topic)
+
     def _handle_async_requests(self, data):
         address_uuid, uuid, response = data
         if uuid not in self._pending_requests:
@@ -932,7 +970,7 @@ def _subscribe(self, alias: str, handlers: Dict[Union[bytes, str], Any]):
         curated_handlers = topics_to_bytes(handlers)
         # Subscribe to topics
         for topic in curated_handlers.keys():
-            self.socket[alias].setsockopt(zmq.SUBSCRIBE, topic)
+            self.subscribe_socket_to_topic(alias, topic)
         # Reset handlers
         self._set_handler(self.socket[alias], curated_handlers)