Add graceful shutdown logic for WebSocket

Signed-off-by: JmPotato <[email protected]>

Author: JmPotato

hyperliquid/info.py

@@ -526,3 +526,8 @@ def unsubscribe(self, subscription: Subscription, subscription_id: int) -> bool:
 
     def name_to_asset(self, name: str) -> int:
         return self.coin_to_asset[self.name_to_coin[name]]
+
+    def close(self):
+        if not hasattr(self, "ws_manager") or self.ws_manager is None:
+            return
+        self.ws_manager.stop()

hyperliquid/websocket_manager.py

@@ -71,14 +71,15 @@ def __init__(self, base_url):
         self.ping_sender = threading.Thread(target=self.send_ping)
 
     def run(self):
-        self.ping_sender.start()
         self.ws.run_forever()
+        self.ping_sender.start()
 
     def send_ping(self):
-        while True:
+        while self.ws.run_forever:
             time.sleep(50)
             logging.debug("Websocket sending ping")
             self.ws.send(json.dumps({"method": "ping"}))
+        logging.debug("Websocket ping sender stopped")
 
     def on_message(self, _ws, message):
         if message == "Websocket connection established.":
@@ -136,3 +137,8 @@ def unsubscribe(self, subscription: Subscription, subscription_id: int) -> bool:
             self.ws.send(json.dumps({"method": "unsubscribe", "subscription": subscription}))
         self.active_subscriptions[identifier] = new_active_subscriptions
         return len(active_subscriptions) != len(new_active_subscriptions)
+
+    def stop(self):
+        if not self.ws.keep_running:
+            return
+        self.ws.close()