Fix Polygon handler registry and add symbol specification

Author: ttt733

alpaca_trade_api/polygon/stream2.py

@@ -17,6 +17,7 @@ def __init__(self, key_id=None):
             'wss://alpaca.socket.polygon.io/stocks'
         ).rstrip('/')
         self._handlers = {}
+        self._handler_symbols = {}
         self._ws = None
         self._retry = int(os.environ.get('APCA_RETRY_MAX', 3))
         self._retry_wait = int(os.environ.get('APCA_RETRY_WAIT', 3))
@@ -114,7 +115,7 @@ def run(self, initial_channels=[]):
             loop.run_until_complete(self.close())
 
     async def close(self):
-        '''Close any of open connections'''
+        '''Close any open connections'''
         if self._ws is not None:
             await self._ws.close()
 
@@ -168,15 +169,10 @@ def _cast(self, subject, data):
     async def _dispatch(self, channel, msg):
         for pat, handler in self._handlers.items():
             if pat.match(channel):
-                ent = self._cast(channel, msg)
-                await handler(self, channel, ent)
-
-    def on(self, channel_pat):
-        def decorator(func):
-            self.register(channel_pat, func)
-            return func
-
-        return decorator
+                handled_symbols = self._handler_symbols.get(handler)
+                if handled_symbols is None or msg['sym'] in handled_symbols:
+                    ent = self._cast(channel, msg)
+                    await handler(self, channel, ent)
 
     def register(self, channel_pat, func):
         if not asyncio.iscoroutinefunction(func):
@@ -188,4 +184,6 @@ def register(self, channel_pat, func):
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
+        handler = self._handlers[channel_pat]
+        del self._handler_symbols[handler]
         del self._handlers[channel_pat]

alpaca_trade_api/stream2.py

@@ -13,6 +13,7 @@ def __init__(self, key_id=None, secret_key=None, base_url=None):
         base_url = re.sub(r'^http', 'ws', base_url or get_base_url())
         self._endpoint = base_url + '/stream'
         self._handlers = {}
+        self._handler_symbols = {}
         self._base_url = base_url
         self._ws = None
         self.polygon = None
@@ -65,7 +66,8 @@ async def _ensure_polygon(self):
         if 'staging' in self._base_url:
             key_id += '-staging'
         self.polygon = polygon.StreamConn(key_id)
-        self.polygon.register(r'.*', self._dispatch_polygon)
+        self.polygon._handlers = self._handlers
+        self.polygon._handler_symbols = self._handler_symbols
         await self.polygon.connect()
 
     async def _ensure_ws(self):
@@ -121,20 +123,17 @@ def _cast(self, channel, msg):
             return Account(msg)
         return Entity(msg)
 
-    async def _dispatch_polygon(self, conn, subject, data):
-        for pat, handler in self._handlers.items():
-            if pat.match(subject):
-                await handler(self, subject, data)
-
     async def _dispatch(self, channel, msg):
         for pat, handler in self._handlers.items():
             if pat.match(channel):
                 ent = self._cast(channel, msg['data'])
                 await handler(self, channel, ent)
 
-    def on(self, channel_pat):
+    def on(self, channel_pat, symbols=None):
         def decorator(func):
             self.register(channel_pat, func)
+            if symbols:
+                self._handler_symbols[func] = symbols
             return func
 
         return decorator
@@ -145,8 +144,14 @@ def register(self, channel_pat, func):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
         self._handlers[channel_pat] = func
+        if self.polygon:
+            self.polygon.register(channel_pat, func)
 
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
+        handler = self._handlers[channel_pat]
+        del self._handler_symbols[handler]
         del self._handlers[channel_pat]
+        if self.polygon:
+            self.polygon.deregister(channel_pat)