diff --git a/examples/websocket/latency.py b/examples/websocket/latency.py new file mode 100644 index 00000000..57a0a183 --- /dev/null +++ b/examples/websocket/latency.py @@ -0,0 +1,17 @@ +from polygon import WebSocketClient +from polygon.websocket.models import WebSocketMessage, EquityQuote +from typing import List, cast +import time + +c = WebSocketClient(subscriptions=["Q.SPY"]) + + +def handle_msg(msgs: List[WebSocketMessage]): + for m in msgs: + q: EquityQuote = cast(EquityQuote, m) + if q.timestamp is not None: + now = time.time() * 1000 + print(now, q.timestamp, now - q.timestamp) + + +c.run(handle_msg) diff --git a/polygon/websocket/__init__.py b/polygon/websocket/__init__.py index 21f6bfcb..66a9dd9a 100644 --- a/polygon/websocket/__init__.py +++ b/polygon/websocket/__init__.py @@ -118,16 +118,20 @@ async def connect( self.subs = set(self.scheduled_subs) self.schedule_resub = False - cmsg: Union[ - List[WebSocketMessage], Union[str, bytes] - ] = await s.recv() - # we know cmsg is Data - msgJson = json.loads(cmsg) # type: ignore - for m in msgJson: - if m["ev"] == "status": - logger.debug("status: %s", m["message"]) - continue + try: + cmsg: Union[ + List[WebSocketMessage], Union[str, bytes] + ] = await asyncio.wait_for(s.recv(), timeout=1) + except asyncio.TimeoutError: + continue + if not self.raw: + # we know cmsg is Data + msgJson = json.loads(cmsg) # type: ignore + for m in msgJson: + if m["ev"] == "status": + logger.debug("status: %s", m["message"]) + continue cmsg = parse(msgJson, logger) if len(cmsg) > 0: