Handle trade error cancels and corrections as auto subscribed messages with trades (#536)

Author: ccnlui

alpaca_trade_api/entity_v2.py

@@ -61,6 +61,32 @@
     "z": "tape"
 }
 
+cancel_error_mapping_v2 = {
+    "S": "symbol",
+    "i": "id",
+    "x": "exchange",
+    "p": "price",
+    "s": "size",
+    "a": "cancel_error_action",
+    "z": "tape",
+    "t": "timestamp",
+}
+
+correction_mapping_v2 = {
+    "S": "symbol",
+    "x": "exchange",
+    "oi": "original_id",
+    "op": "original_price",
+    "os": "original_size",
+    "oc": "original_conditions",
+    "ci": "corrected_id",
+    "cp": "corrected_price",
+    "cs": "corrected_size",
+    "cc": "corrected_conditions",
+    "z": "tape",
+    "t": "timestamp",
+}
+
 
 class EntityListType(Enum):
     Trade = Trade, trade_mapping_v2
@@ -152,6 +178,20 @@ def __init__(self, raw):
         super().__init__(luld_mapping_v2, raw)
 
 
+class CancelErrorV2(Remapped, _NanoTimestamped, Entity):
+    _tskeys = ('t',)
+
+    def __init__(self, raw):
+        super().__init__(cancel_error_mapping_v2, raw)
+
+
+class CorrectionV2(Remapped, _NanoTimestamped, Entity):
+    _tskeys = ('t',)
+
+    def __init__(self, raw):
+        super().__init__(correction_mapping_v2, raw)
+
+
 class SnapshotV2:
     def __init__(self, raw):
         self.latest_trade = _convert_or_none(TradeV2, raw.get('latestTrade'))

alpaca_trade_api/stream.py

@@ -14,8 +14,22 @@
 
 from .common import get_base_url, get_data_stream_url, get_credentials, URL
 from .entity import Entity
-from .entity_v2 import quote_mapping_v2, trade_mapping_v2, bar_mapping_v2, \
-    status_mapping_v2, luld_mapping_v2, Trade, Quote, Bar, StatusV2, LULDV2
+from .entity_v2 import (
+    quote_mapping_v2,
+    trade_mapping_v2,
+    bar_mapping_v2,
+    status_mapping_v2,
+    luld_mapping_v2,
+    cancel_error_mapping_v2,
+    correction_mapping_v2,
+    Trade,
+    Quote,
+    Bar,
+    StatusV2,
+    LULDV2,
+    CancelErrorV2,
+    CorrectionV2,
+)
 
 log = logging.getLogger(__name__)
 
@@ -167,7 +181,10 @@ def _subscribe(self, handler, symbols, handlers):
             asyncio.get_event_loop().run_until_complete(self._subscribe_all())
 
     async def _subscribe_all(self):
-        if any(self._handlers.values()):
+        if any(
+            v for k, v in self._handlers.items()
+            if k not in ("cancelErrors", "corrections")
+        ):
             msg = {
                 k: tuple(v.keys())
                 for k, v in self._handlers.items()
@@ -193,7 +210,10 @@ async def _unsubscribe(self,
 
     async def _run_forever(self):
         # do not start the websocket connection until we subscribe to something
-        while not any(self._handlers.values()):
+        while not any(
+            v for k, v in self._handlers.items()
+            if k not in ("cancelErrors", "corrections")
+        ):
             if not self._stop_stream_queue.empty():
                 # the ws was signaled to stop before starting the loop so
                 # we break
@@ -283,6 +303,8 @@ def __init__(self,
                          )
         self._handlers['statuses'] = {}
         self._handlers['lulds'] = {}
+        self._handlers['cancelErrors'] = {}
+        self._handlers['corrections'] = {}
         self._name = 'stock data'
 
     def _cast(self, msg_type, msg):
@@ -298,6 +320,16 @@ def _cast(self, msg_type, msg):
                     luld_mapping_v2[k]: v
                     for k, v in msg.items() if k in luld_mapping_v2
                 })
+            elif msg_type == 'x':
+                result = CancelErrorV2({
+                    cancel_error_mapping_v2[k]: v
+                    for k, v in msg.items() if k in cancel_error_mapping_v2
+                })
+            elif msg_type == 'c':
+                result = CorrectionV2({
+                    correction_mapping_v2[k]: v
+                    for k, v in msg.items() if k in correction_mapping_v2
+                })
         return result
 
     async def _dispatch(self, msg):
@@ -313,6 +345,16 @@ async def _dispatch(self, msg):
                 symbol, self._handlers['lulds'].get('*', None))
             if handler:
                 await handler(self._cast(msg_type, msg))
+        elif msg_type == 'x':
+            handler = self._handlers['cancelErrors'].get(
+                symbol, self._handlers['cancelErrors'].get('*', None))
+            if handler:
+                await handler(self._cast(msg_type, msg))
+        elif msg_type == 'c':
+            handler = self._handlers['corrections'].get(
+                symbol, self._handlers['corrections'].get('*', None))
+            if handler:
+                await handler(self._cast(msg_type, msg))
         else:
             await super()._dispatch(msg)
 
@@ -355,6 +397,16 @@ def unsubscribe_lulds(self, *symbols):
         for symbol in symbols:
             del self._handlers['lulds'][symbol]
 
+    def register_handler(self, msg_type, handler, *symbols):
+        if handler is not None:
+            _ensure_coroutine(handler)
+            for symbol in symbols:
+                self._handlers[msg_type][symbol] = handler
+
+    def unregister_handler(self, msg_type, *symbols):
+        for symbol in symbols:
+            del self._handlers[msg_type][symbol]
+
 
 class CryptoDataStream(_DataStream):
     def __init__(self,
@@ -531,8 +583,20 @@ def __init__(self,
     def subscribe_trade_updates(self, handler):
         self._trading_ws.subscribe_trade_updates(handler)
 
-    def subscribe_trades(self, handler, *symbols):
+    def subscribe_trades(
+        self,
+        handler,
+        *symbols,
+        handler_cancel_errors=None,
+        handler_corrections=None
+    ):
         self._data_ws.subscribe_trades(handler, *symbols)
+        self._data_ws.register_handler("cancelErrors",
+                                       handler_cancel_errors,
+                                       *symbols)
+        self._data_ws.register_handler("corrections",
+                                       handler_corrections,
+                                       *symbols)
 
     def subscribe_quotes(self, handler, *symbols):
         self._data_ws.subscribe_quotes(handler, *symbols)
@@ -607,6 +671,20 @@ def decorator(func):
 
         return decorator
 
+    def on_cancel_error(self, *symbols):
+        def decorator(func):
+            self._data_ws.register_handler("cancelErrors", func, *symbols)
+            return func
+
+        return decorator
+
+    def on_corrections(self, *symbols):
+        def decorator(func):
+            self._data_ws.register_handler("corrections", func, *symbols)
+            return func
+
+        return decorator
+
     def on_crypto_trade(self, *symbols):
         def decorator(func):
             self.subscribe_crypto_trades(func, *symbols)
@@ -637,6 +715,8 @@ def decorator(func):
 
     def unsubscribe_trades(self, *symbols):
         self._data_ws.unsubscribe_trades(*symbols)
+        self._data_ws.unregister_handler("cancelErrors", *symbols)
+        self._data_ws.unregister_handler("corrections", *symbols)
 
     def unsubscribe_quotes(self, *symbols):
         self._data_ws.unsubscribe_quotes(*symbols)