Add websocket support

Author: justinpolygon

polygon/websocket/__init__.py

@@ -140,7 +140,7 @@ async def connect(
                             if m["ev"] == "status":
                                 logger.debug("status: %s", m["message"])
                                 continue
-                        cmsg = parse(msgJson, logger)
+                        cmsg = parse(msgJson, logger, self.market)
 
                     if len(cmsg) > 0:
                         await processor(cmsg)  # type: ignore

polygon/websocket/models/__init__.py

@@ -3,47 +3,72 @@
 from .models import *
 import logging
 
+# Define the mapping of market and event type to model class
+MARKET_EVENT_MAP = {
+    Market.Stocks: {
+        "A": EquityAgg,
+        "AM": EquityAgg,
+        "T": EquityTrade,
+        "Q": EquityQuote,
+        "LULD": LimitUpLimitDown,
+        "FMV": FairMarketValue,
+        "NOI": Imbalance,
+        "LV": LaunchpadValue,
+    },
+    Market.Options: {
+        "A": EquityAgg,
+        "AM": EquityAgg,
+        "T": EquityTrade,
+        "Q": EquityQuote,
+        "FMV": FairMarketValue,
+    },
+    Market.Indices: {
+        "A": EquityAgg,
+        "AM": EquityAgg,
+        "V": IndexValue,
+    },
+    Market.Futures: {
+        "A": FuturesAgg,
+        "AM": FuturesAgg,
+        "T": FuturesTrade,
+        "Q": FuturesQuote,
+    },
+    Market.Crypto: {
+        "XA": CurrencyAgg,
+        "XAS": CurrencyAgg,
+        "XT": CryptoTrade,
+        "XQ": CryptoQuote,
+        "XL2": Level2Book,
+        "FMV": FairMarketValue,
+    },
+    Market.Forex: {
+        "CA": CurrencyAgg,
+        "CAS": CurrencyAgg,
+        "C": ForexQuote,
+        "FMV": FairMarketValue,
+    },
+}
 
-def parse_single(data: Dict[str, Any]):
+
+def parse_single(data: Dict[str, Any], market: Market) -> Any:
     event_type = data["ev"]
-    if event_type in [EventType.EquityAgg.value, EventType.EquityAggMin.value]:
-        return EquityAgg.from_dict(data)
-    elif event_type in [
-        EventType.CryptoAgg.value,
-        EventType.CryptoAggSec.value,
-        EventType.ForexAgg.value,
-        EventType.ForexAggSec.value,
-    ]:
-        return CurrencyAgg.from_dict(data)
-    elif event_type == EventType.EquityTrade.value:
-        return EquityTrade.from_dict(data)
-    elif event_type == EventType.CryptoTrade.value:
-        return CryptoTrade.from_dict(data)
-    elif event_type == EventType.EquityQuote.value:
-        return EquityQuote.from_dict(data)
-    elif event_type == EventType.ForexQuote.value:
-        return ForexQuote.from_dict(data)
-    elif event_type == EventType.CryptoQuote.value:
-        return CryptoQuote.from_dict(data)
-    elif event_type == EventType.Imbalances.value:
-        return Imbalance.from_dict(data)
-    elif event_type == EventType.LimitUpLimitDown.value:
-        return LimitUpLimitDown.from_dict(data)
-    elif event_type == EventType.CryptoL2.value:
-        return Level2Book.from_dict(data)
-    elif event_type == EventType.Value.value:
-        return IndexValue.from_dict(data)
-    elif event_type == EventType.LaunchpadValue.value:
-        return LaunchpadValue.from_dict(data)
-    elif event_type == EventType.BusinessFairMarketValue.value:
-        return FairMarketValue.from_dict(data)
-    return None
+    # Look up the model class based on market and event type
+    model_class = MARKET_EVENT_MAP.get(market, {}).get(event_type)
+    if model_class:
+        return model_class.from_dict(data)
+    else:
+        # Log a warning for unrecognized event types, unless it's a status message
+        if event_type != "status":
+            logger.warning("Unknown event type '%s' for market %s", event_type, market)
+        return None
 
 
-def parse(msg: List[Dict[str, Any]], logger: logging.Logger) -> List[WebSocketMessage]:
+def parse(
+    msg: List[Dict[str, Any]], logger: logging.Logger, market: Market
+) -> List[WebSocketMessage]:
     res = []
     for m in msg:
-        parsed = parse_single(m)
+        parsed = parse_single(m, market)
         if parsed is None:
             if m["ev"] != "status":
                 logger.warning("could not parse message %s", m)

polygon/websocket/models/common.py

@@ -11,7 +11,6 @@ class Feed(Enum):
     Launchpad = "launchpad.polygon.io"
     Business = "business.polygon.io"
     EdgxBusiness = "edgx-business.polygon.io"
-    IEXBusiness = "iex-business.polygon.io"
     DelayedBusiness = "delayed-business.polygon.io"
     DelayedEdgxBusiness = "delayed-edgx-business.polygon.io"
     DelayedNasdaqLastSaleBusiness = "delayed-nasdaq-last-sale-business.polygon.io"
@@ -28,6 +27,7 @@ class Market(Enum):
     Forex = "forex"
     Crypto = "crypto"
     Indices = "indices"
+    Futures = "futures"
 
 
 class EventType(Enum):
@@ -46,12 +46,10 @@ class EventType(Enum):
     LimitUpLimitDown = "LULD"
     CryptoL2 = "XL2"
     Value = "V"
-    """Launchpad* EventTypes are only available to Launchpad users. These values are the same across all asset classes (
-    stocks, options, forex, crypto).
-    """
     LaunchpadValue = "LV"
     LaunchpadAggMin = "AM"
-    """Business* EventTypes are only available to Business users. These values are the same across all asset classes (
-    stocks, options, forex, crypto).
-    """
     BusinessFairMarketValue = "FMV"
+    FuturesTrade = "T"
+    FuturesQuote = "Q"
+    FuturesAgg = "A"
+    FuturesAggMin = "AM"

polygon/websocket/models/models.py

@@ -359,6 +359,93 @@ def from_dict(d):
         )
 
 
+@modelclass
+class FuturesTrade:
+    event_type: Optional[str] = None
+    symbol: Optional[str] = None
+    price: Optional[float] = None
+    size: Optional[int] = None
+    timestamp: Optional[int] = None
+    sequence_number: Optional[int] = None
+
+    @staticmethod
+    def from_dict(d):
+        return FuturesTrade(
+            event_type=d.get("ev"),
+            symbol=d.get("sym"),
+            price=d.get("p"),
+            size=d.get("s"),
+            timestamp=d.get("t"),
+            sequence_number=d.get("q"),
+        )
+
+
+@modelclass
+class FuturesQuote:
+    event_type: Optional[str] = None
+    symbol: Optional[str] = None
+    bid_price: Optional[float] = None
+    bid_size: Optional[int] = None
+    bid_timestamp: Optional[int] = None
+    ask_price: Optional[float] = None
+    ask_size: Optional[int] = None
+    ask_timestamp: Optional[int] = None
+    sip_timestamp: Optional[int] = None
+
+    @staticmethod
+    def from_dict(d):
+        return FuturesQuote(
+            event_type=d.get("ev"),
+            symbol=d.get("sym"),
+            bid_price=d.get("bp"),
+            bid_size=d.get("bs"),
+            bid_timestamp=d.get("bt"),
+            ask_price=d.get("ap"),
+            ask_size=d.get("as"),
+            ask_timestamp=d.get("at"),
+            sip_timestamp=d.get("t"),
+        )
+
+
+@modelclass
+class FuturesAgg:
+    event_type: Optional[str] = None
+    symbol: Optional[str] = None
+    volume: Optional[float] = None
+    accumulated_volume: Optional[float] = None
+    official_open_price: Optional[float] = None
+    vwap: Optional[float] = None
+    open: Optional[float] = None
+    close: Optional[float] = None
+    high: Optional[float] = None
+    low: Optional[float] = None
+    aggregate_vwap: Optional[float] = None
+    average_size: Optional[float] = None
+    start_timestamp: Optional[int] = None
+    end_timestamp: Optional[int] = None
+    otc: Optional[bool] = None  # If present
+
+    @staticmethod
+    def from_dict(d):
+        return FuturesAgg(
+            event_type=d.get("ev"),
+            symbol=d.get("sym"),
+            volume=d.get("v"),
+            accumulated_volume=d.get("av"),
+            official_open_price=d.get("op"),
+            vwap=d.get("vw"),
+            open=d.get("o"),
+            close=d.get("c"),
+            high=d.get("h"),
+            low=d.get("l"),
+            aggregate_vwap=d.get("a"),
+            average_size=d.get("z"),
+            start_timestamp=d.get("s"),
+            end_timestamp=d.get("e"),
+            otc=d.get("otc"),
+        )
+
+
 WebSocketMessage = NewType(
     "WebSocketMessage",
     List[
@@ -376,6 +463,9 @@ def from_dict(d):
             IndexValue,
             LaunchpadValue,
             FairMarketValue,
+            FuturesTrade,
+            FuturesQuote,
+            FuturesAgg,
         ]
     ],
 )