stream crypto orderbook (#597)

* stream crypto orderbook

* flake8 format

* sonarcloud check

Author: ccnlui

alpaca_trade_api/entity_v2.py

@@ -87,6 +87,14 @@
     "t": "timestamp",
 }
 
+orderbook_mapping_v2 = {
+    "S": "symbol",
+    "x": "exchange",
+    "t": "timestamp",
+    "b": "bids",
+    "a": "asks",
+}
+
 
 class EntityListType(Enum):
     Trade = Trade, trade_mapping_v2
@@ -225,6 +233,22 @@ def __init__(self, raw):
             self[k] = _convert_or_none(QuoteV2, v)
 
 
+class BidOrAsk(Entity):
+    def __init__(self, raw):
+        super().__init__(raw)
+
+
+class OrderbookV2(Entity):
+    def __init__(self, raw):
+        super().__init__(raw)
+        if self.bids:
+            for i in range(len(self.bids)):
+                self.bids[i] = BidOrAsk(self.bids[i])
+        if self.asks:
+            for i in range(len(self.asks)):
+                self.asks[i] = BidOrAsk(self.asks[i])
+
+
 class NewsV2(Entity):
     def __init__(self, raw):
         super().__init__(raw)

alpaca_trade_api/stream.py

@@ -18,6 +18,7 @@
     luld_mapping_v2,
     cancel_error_mapping_v2,
     correction_mapping_v2,
+    orderbook_mapping_v2,
     Trade,
     Quote,
     Bar,
@@ -26,6 +27,7 @@
     CancelErrorV2,
     CorrectionV2,
     NewsV2,
+    OrderbookV2,
 )
 
 log = logging.getLogger(__name__)
@@ -219,16 +221,7 @@ async def _unsubscribe(self,
                            bars=(),
                            updated_bars=(),
                            daily_bars=()):
-        if trades or quotes or bars or updated_bars or daily_bars:
-            await self._ws.send(
-                msgpack.packb({
-                    'action':      'unsubscribe',
-                    'trades':      trades,
-                    'quotes':      quotes,
-                    'bars':        bars,
-                    'updatedBars': updated_bars,
-                    'dailyBars':   daily_bars,
-                }))
+        raise NotImplementedError()
 
     async def _run_forever(self):
         self._loop = asyncio.get_running_loop()
@@ -311,8 +304,9 @@ def unsubscribe_bars(self, *symbols):
 
     def unsubscribe_updated_bars(self, *symbols):
         if self._running:
-            asyncio.get_event_loop().run_until_complete(
-                self._unsubscribe(updated_bars=symbols))
+            asyncio.run_coroutine_threadsafe(
+                self._unsubscribe(updated_bars=symbols),
+                self._loop).result()
         for symbol in symbols:
             del self._handlers['updatedBars'][symbol]
 
@@ -482,8 +476,63 @@ def __init__(self,
                          raw_data=raw_data,
                          websocket_params=websocket_params,
                          )
+        self._handlers['orderbooks'] = {}
         self._name = 'crypto data'
 
+    def _cast(self, msg_type, msg):
+        result = super()._cast(msg_type, msg)
+        if not self._raw_data:
+            if msg_type == 'o':
+                result = OrderbookV2({
+                    orderbook_mapping_v2[k]: v
+                    for k, v in msg.items() if k in orderbook_mapping_v2
+                })
+        return result
+
+    async def _dispatch(self, msg):
+        msg_type = msg.get('T')
+        symbol = msg.get('S')
+        if msg_type == 'o':
+            handler = self._handlers['orderbooks'].get(
+                symbol, self._handlers['orderbooks'].get('*', None))
+            if handler:
+                await handler(self._cast(msg_type, msg))
+        else:
+            await super()._dispatch(msg)
+
+    async def _unsubscribe(self,
+                           trades=(),
+                           quotes=(),
+                           orderbooks=(),
+                           bars=(),
+                           updated_bars=(),
+                           daily_bars=()):
+        if (
+            trades or quotes or orderbooks or bars or updated_bars
+            or daily_bars
+        ):
+            await self._ws.send(
+                msgpack.packb({
+                    'action':      'unsubscribe',
+                    'trades':      trades,
+                    'quotes':      quotes,
+                    'orderbooks':  orderbooks,
+                    'bars':        bars,
+                    'updatedBars': updated_bars,
+                    'dailyBars':   daily_bars,
+                }))
+
+    def subscribe_orderbooks(self, handler, *symbols):
+        self._subscribe(handler, symbols, self._handlers['orderbooks'])
+
+    def unsubscribe_orderbooks(self, *symbols):
+        if self._running:
+            asyncio.run_coroutine_threadsafe(
+                self._unsubscribe(orderbooks=symbols),
+                self._loop).result()
+        for symbol in symbols:
+            del self._handlers['orderbooks'][symbol]
+
 
 class NewsDataStream(_DataStream):
     def __init__(self,
@@ -529,7 +578,7 @@ async def _unsubscribe(self, news=()):
         if news:
             await self._ws.send(
                 msgpack.packb({
-                    'action':    'unsubscribe',
+                    'action': 'unsubscribe',
                     'news':    news,
                 }))
 
@@ -779,6 +828,9 @@ def subscribe_crypto_updated_bars(self, handler, *symbols):
     def subscribe_crypto_daily_bars(self, handler, *symbols):
         self._crypto_ws.subscribe_daily_bars(handler, *symbols)
 
+    def subscribe_crypto_orderbooks(self, handler, *symbols):
+        self._crypto_ws.subscribe_orderbooks(handler, *symbols)
+
     def subscribe_news(self, handler, *symbols):
         self._news_ws.subscribe_news(handler, *symbols)
 
@@ -842,7 +894,7 @@ def decorator(func):
 
         return decorator
 
-    def on_corrections(self, *symbols):
+    def on_correction(self, *symbols):
         def decorator(func):
             self._data_ws.register_handler("corrections", func, *symbols)
             return func
@@ -884,6 +936,13 @@ def decorator(func):
 
         return decorator
 
+    def on_crypto_orderbook(self, *symbols):
+        def decorator(func):
+            self.subscribe_crypto_orderbooks(func, *symbols)
+            return func
+
+        return decorator
+
     def on_news(self, *symbols):
         def decorator(func):
             self.subscribe_news(func, *symbols)
@@ -929,6 +988,9 @@ def unsubscribe_crypto_updated_bars(self, *symbols):
     def unsubscribe_crypto_daily_bars(self, *symbols):
         self._crypto_ws.unsubscribe_daily_bars(*symbols)
 
+    def unsubscribe_crypto_orderbooks(self, *symbols):
+        self._crypto_ws.unsubscribe_orderbooks(*symbols)
+
     def unsubscribe_news(self, *symbols):
         self._news_ws.unsubscribe_news(*symbols)