Merge pull request #90 from alpacahq/feature/ws-handler-registry

Fix Polygon handler registry and add symbol specification

Author: ttt733

alpaca_trade_api/polygon/stream2.py

@@ -18,6 +18,7 @@ def __init__(self, key_id=None):
             'wss://alpaca.socket.polygon.io/stocks'
         ).rstrip('/')
         self._handlers = {}
+        self._handler_symbols = {}
         self._streams = set([])
         self._ws = None
         self._retry = int(os.environ.get('APCA_RETRY_MAX', 3))
@@ -180,7 +181,7 @@ def run(self, initial_channels=[]):
             loop.close()
 
     async def close(self):
-        '''Close any of open connections'''
+        '''Close any open connections'''
         if self._ws is not None:
             await self._ws.close()
         self._ws = None
@@ -236,24 +237,21 @@ async def _dispatch(self, msg):
         channel = msg.get('ev')
         for pat, handler in self._handlers.items():
             if pat.match(channel):
-                ent = self._cast(channel, msg)
-                await handler(self, channel, ent)
+                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 on(self, channel_pat):
-        def decorator(func):
-            self.register(channel_pat, func)
-            return func
-
-        return decorator
-
-    def register(self, channel_pat, func):
+    def register(self, channel_pat, func, symbols=None):
         if not asyncio.iscoroutinefunction(func):
             raise ValueError('handler must be a coroutine function')
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
         self._handlers[channel_pat] = func
+        self._handler_symbols[func] = symbols
 
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
+        self._handler_symbols.pop(self._handlers[channel_pat], None)
         del self._handlers[channel_pat]

alpaca_trade_api/stream2.py

@@ -14,6 +14,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
@@ -67,7 +68,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):
@@ -148,32 +150,33 @@ 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)
+            self.register(channel_pat, func, symbols)
             return func
 
         return decorator
 
-    def register(self, channel_pat, func):
+    def register(self, channel_pat, func, symbols=None):
         if not asyncio.iscoroutinefunction(func):
             raise ValueError('handler must be a coroutine function')
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
         self._handlers[channel_pat] = func
+        self._handler_symbols[func] = symbols
+        if self.polygon:
+            self.polygon.register(channel_pat, func, symbols)
 
     def deregister(self, channel_pat):
         if isinstance(channel_pat, str):
             channel_pat = re.compile(channel_pat)
+        self._handler_symbols.pop(self._handlers[channel_pat], None)
         del self._handlers[channel_pat]
+        if self.polygon:
+            self.polygon.deregister(channel_pat)

tests/test_stream2.py

@@ -1,6 +1,6 @@
 from alpaca_trade_api.stream2 import StreamConn
+from alpaca_trade_api.polygon.stream2 import StreamConn as PolyStream
 from alpaca_trade_api.entity import Account
-from alpaca_trade_api.polygon.entity import Entity as PolyEntity
 import asyncio
 import json
 
@@ -126,12 +126,15 @@ async def on_raise(conn, stream, msg):
     ent = conn._cast('other', {'key': 'value'})
     assert ent.key == 'value'
 
-    # _dispatch_nats
+    # polygon _dispatch
     conn = StreamConn('key-id', 'secret-key')
+    conn.polygon = PolyStream('key-id')
+    msg_data = {'key': 'value', 'ev': 'Q'}
+    conn.polygon._cast = mock.Mock(return_value=msg_data)
 
-    @conn.on('^Q.')
+    @conn.on('Q')
     async def on_q(conn, subject, data):
         on_q.data = data
 
-    _run(conn._dispatch_polygon(conn, 'Q.SPY', PolyEntity({'key': 'value'})))
-    assert on_q.data.key == 'value'
+    _run(conn.polygon._dispatch(msg_data))
+    assert on_q.data['key'] == 'value'