Fixes by at-cf (#297)

* Fixes by at-cf

* fix lint

Author: clickingbuttons

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)

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: