add unsubscribe mechanism

Author: timothy.tamm

.gitignore

@@ -0,0 +1,2 @@
+*.pyc
+*/__pycache__

comm/asyncClient.py

@@ -33,7 +33,7 @@ def connect(self):
     def connection_made(self, transport):
         super().connection_made(transport)
         if len(self.subscriptions) > 0:
-            self.subscribe(self.subscriptions)
+            self.subscribe(self.subscriptions, Subscribe.SUBSCRIBE)
 
     def msg_received(self, data, msg_type):
         msg = message_buffers[msg_type]()
@@ -43,10 +43,11 @@ def msg_received(self, data, msg_type):
     def write(self, msg, msg_type):
         super().write(msg, msg_type)
 
-    def subscribe(self, msg_types):
+    def subscribe(self, msg_types, direction):
         msg = Subscribe()
         for msg_type in msg_types:
             msg.msg_types.append(msg_type.value)
+        msg.dir = direction
         self.write(msg.SerializeToString(), MsgType.SUBSCRIBE)
 
 

comm/subscribe.proto

@@ -4,4 +4,10 @@ package mateROV;
 
 message Subscribe {
   repeated int32 msg_types = 1;
+  required Direction dir = 2;
+
+  enum Direction {
+    SUBSCRIBE = 0;
+    UNSUBSCRIBE = 1;
+  }
 }

comm/subscribe_pb2.py

@@ -14,6 +14,8 @@ class MockGuiModule(ProtoModule):
     def __init__(self, loop):
         super().__init__(loop, ADDRESS, PORT,[MsgType.MOCK_MSG])
         self.value = -1
+        self.sub_ticks = 0
+        self.subbed = True
 
     def msg_received(self, msg, msg_type):
         # This gets called whenever any message is received
@@ -27,6 +29,19 @@ def tick(self):
         # for this mock module we will print out the current value
         print('Current value: {}'.format(self.value))
 
+        # to demonstrate subscription and unsubscription,
+        # we will periodically unsubscribe and resubscribe
+        if self.sub_ticks > 100:
+            if self.subbed:
+                print('Unsubscribed!')
+                self.unsubscribe([MsgType.MOCK_MSG])
+            else:
+                print('Subscribed!')
+                self.subscribe([MsgType.MOCK_MSG])
+            self.subbed = not self.subbed
+            self.sub_ticks = 0
+        self.sub_ticks += 1
+
 
 def main():
     loop = asyncio.get_event_loop()

mockGuiModule.py

@@ -1,4 +1,5 @@
 from comm.asyncClient import AsyncClient
+from comm.subscribe_pb2 import Subscribe
 
 class ProtoModule:
     def __init__(self, loop, addr, port, subscriptions):
@@ -13,7 +14,10 @@ def msg_received(self, msg, msg_type):
         raise NotImplementedError()
 
     def subscribe(self, msg_types):
-        self.client.subscribe(msg_types)
+        self.client.subscribe(msg_types, Subscribe.SUBSCRIBE)
+
+    def unsubscribe(self, msg_types):
+        self.client.subscribe(msg_types, Subscribe.UNSUBSCRIBE)
 
     def write(self, msg, msg_type):
         self.client.write(msg, msg_type)

modules/protoModule.py

@@ -3,6 +3,7 @@
 import time, os, sys, logging
 import asyncio # minimum Python 3.4, changed in 3.5.1
 from comm.serverProto import ServerProto
+from comm.subscribe_pb2 import Subscribe
 from comm.constants import *
 
 ADDRESS = os.environ.get("BIND_ADDRESS","localhost")
@@ -17,12 +18,25 @@ def __init__(self, loop):
         coro = loop.create_server(lambda: ServerProto(self), ADDRESS, PORT)
         self.server = loop.run_until_complete(coro)
 
+    def _handle_subscriptions(self, protocol, data):
+        if data.dir == Subscribe.SUBSCRIBE:
+            self._add_subscriptions(protocol, data)
+        else:
+            self._remove_subscriptions(protocol, data)
+
+    def _remove_subscriptions(self, protocol, data):
+        for msg_type in data.msg_types:
+            m_type = MsgType(msg_type)
+            if m_type in self.subs:
+                self.subs[m_type].remove(protocol)
+
     def _add_subscriptions(self, protocol, data):
         for msg_type in data.msg_types:
-            if msg_type in self.subs:
-                self.subs[MsgType(msg_type)].append(protocol)
+            m_type = MsgType(msg_type)
+            if m_type in self.subs:
+                self.subs[m_type].append(protocol)
             else:
-                self.subs[MsgType(msg_type)] = [protocol]
+                self.subs[m_type] = [protocol]
 
     def _forward_msg(self, msg, msg_type):
         if msg_type in self.subs:
@@ -34,7 +48,7 @@ def msg_received(self, protocol, msg, msg_type):
         if msg_type == MsgType.SUBSCRIBE:
             data = message_buffers[msg_type]()
             data.ParseFromString(msg)
-            self._add_subscriptions(protocol, data)
+            self._handle_subscriptions(protocol, data)
         else:
             self._forward_msg(msg, msg_type)