resub on reconnect

Author: clickingbuttons

examples/websocket/simple.py

@@ -2,7 +2,7 @@
 from polygon.websocket.models import WebSocketMessage
 from typing import List
 
-c = WebSocketClient(subscriptions=["T.*"])
+c = WebSocketClient(subscriptions=["T.*"], verbose=True)
 
 
 def handle_msg(msgs: List[WebSocketMessage]):

polygon/websocket/__init__.py

@@ -63,7 +63,7 @@ def __init__(
         self.websocket: Optional[WebSocketClientProtocol] = None
         if subscriptions is None:
             subscriptions = []
-        self.scheduled_subs = set(subscriptions)
+        self.scheduled_subs: Set[str] = set(subscriptions)
         self.schedule_resub = True
 
     # https://websockets.readthedocs.io/en/stable/reference/client.html#opening-a-connection
@@ -138,6 +138,9 @@ async def connect(
             except ConnectionClosedError:
                 logger.debug("connection closed (ERR)")
                 reconnects += 1
+                self.scheduled_subs = set(self.subs)
+                self.subs = set()
+                self.schedule_resub = True
                 if self.max_reconnects is not None and reconnects > self.max_reconnects:
                     return
                 continue