wiz-trader 0.9.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

wiz_trader/quotes/client.py

@@ -3,7 +3,7 @@ import json
 import os
 import logging
 import random
-from typing import Callable, List, Optional, Any, Iterator
+from typing import Callable, List, Optional, Any, Dict, Iterator
 
 import websockets
 from websockets.exceptions import ConnectionClosed
@@ -25,10 +25,13 @@ class QuotesClient:
   Attributes:
     base_url (str): WebSocket URL of the quotes server.
     token (str): JWT token for authentication.
-    on_tick (Callable[[dict], None]): Callback to process received tick data.
     log_level (str): Logging level. Options: "error", "info", "debug".
   """
 
+  # Constants for actions
+  ACTION_SUBSCRIBE = "subscribe"
+  ACTION_UNSUBSCRIBE = "unsubscribe"
+
   def __init__(
     self, 
     base_url: Optional[str] = None, 
@@ -57,14 +60,21 @@ class QuotesClient:
     # Construct the WebSocket URL.
     self.url = f"{self.base_url}?token={self.token}"
     self.ws: Optional[websockets.WebSocketClientProtocol] = None
-    self.on_tick: Optional[Callable[[dict], None]] = None
     self.subscribed_instruments: set = set()
+    self._running = False
+    self._background_task = None
 
     # Backoff configuration for reconnection (in seconds)
     self._backoff_base = 1
     self._backoff_factor = 2
     self._backoff_max = 60
 
+    # Callbacks
+    self.on_tick: Optional[Callable[[Any, dict], None]] = None
+    self.on_connect: Optional[Callable[[Any], None]] = None
+    self.on_close: Optional[Callable[[Any, Optional[int], Optional[str]], None]] = None
+    self.on_error: Optional[Callable[[Any, Exception], None]] = None
+    
     logger.debug("Initialized QuotesClient with URL: %s", self.url)
 
   def _chunk_list(self, data: List[Any], chunk_size: int) -> Iterator[List[Any]]:
@@ -81,27 +91,33 @@ class QuotesClient:
     for i in range(0, len(data), chunk_size):
         yield data[i:i + chunk_size]
 
-  async def connect(self) -> None:
+  async def _connect_with_backoff(self) -> None:
     """
-    Continuously connect to the quotes server and process incoming messages.
-    Implements an exponential backoff reconnection strategy.
+    Continuously connect to the quotes server with exponential backoff.
     """
     backoff = self._backoff_base
 
-    while True:
+    while self._running:
       try:
         logger.info("Connecting to %s ...", self.url)
         async with websockets.connect(self.url, max_size=self.max_message_size) as websocket:
           self.ws = websocket
           logger.info("Connected to the quotes server.")
           
+          # Call the on_connect callback if provided
+          if self.on_connect:
+            try:
+              self.on_connect(self)
+            except Exception as e:
+              logger.error("Error in on_connect callback: %s", e, exc_info=True)
+          
           # On reconnection, re-subscribe if needed.
           if self.subscribed_instruments:
             # Split into batches to avoid message size issues
             instruments_list = list(self.subscribed_instruments)
             for batch in self._chunk_list(instruments_list, self.batch_size):
               subscribe_msg = {
-                "action": "subscribe",
+                "action": self.ACTION_SUBSCRIBE,
                 "instruments": batch
               }
               await self.ws.send(json.dumps(subscribe_msg))
@@ -115,9 +131,25 @@ class QuotesClient:
           await self._handle_messages()
       except ConnectionClosed as e:
         logger.info("Disconnected from the quotes server: %s", e)
+        # Call the on_close callback if provided
+        if self.on_close:
+          try:
+            self.on_close(self, getattr(e, 'code', None), str(e))
+          except Exception as e:
+            logger.error("Error in on_close callback: %s", e, exc_info=True)
       except Exception as e:
         logger.error("Connection error: %s", e, exc_info=True)
+        # Call the on_error callback if provided
+        if self.on_error:
+          try:
+            self.on_error(self, e)
+          except Exception as e:
+            logger.error("Error in on_error callback: %s", e, exc_info=True)
 
+      # Don't reconnect if we're no longer running
+      if not self._running:
+        break
+        
       # Exponential backoff before reconnecting.
       sleep_time = min(backoff, self._backoff_max)
       logger.info("Reconnecting in %s seconds...", sleep_time)
@@ -152,7 +184,7 @@ class QuotesClient:
               try:
                 tick = json.loads(json_str)
                 if self.on_tick:
-                  self.on_tick(tick)
+                  self.on_tick(self, tick)
               except json.JSONDecodeError as e:
                 logger.error("Failed to parse JSON object: %s", str(e))
                 logger.error("Invalid JSON: %s...", json_str[:100])
@@ -161,7 +193,7 @@ class QuotesClient:
             try:
               tick = json.loads(message)
               if self.on_tick:
-                self.on_tick(tick)
+                self.on_tick(self, tick)
             except json.JSONDecodeError as e:
               logger.error("Failed to parse JSON: %s", str(e))
               logger.error("Invalid JSON message: %s...", message[:100])
@@ -169,13 +201,19 @@ class QuotesClient:
           logger.warning("Received non-string message: %s", type(message))
     except ConnectionClosed as e:
       logger.info("Connection closed during message handling: %s", e)
+      # Let the _connect_with_backoff method handle reconnection
     except Exception as e:
       logger.error("Error processing message: %s", str(e), exc_info=True)
+      # Call the on_error callback if provided
+      if self.on_error:
+        try:
+          self.on_error(self, e)
+        except Exception as e:
+          logger.error("Error in on_error callback: %s", e, exc_info=True)
 
   async def subscribe(self, instruments: List[str]) -> None:
     """
-    Subscribe to a list of instruments and update the subscription list.
-    Splits large subscription requests into batches to avoid message size issues.
+    Subscribe to a list of instruments.
 
     Args:
       instruments (List[str]): List of instrument identifiers.
@@ -189,7 +227,7 @@ class QuotesClient:
         new_instruments_list = list(new_instruments)
         for batch in self._chunk_list(new_instruments_list, self.batch_size):
           logger.info("Subscribing to batch of %d instruments", len(batch))
-          message = {"action": "subscribe", "instruments": batch}
+          message = {"action": self.ACTION_SUBSCRIBE, "instruments": batch}
           await self.ws.send(json.dumps(message))
           # Small delay between batches to avoid overwhelming the server
           await asyncio.sleep(0.1)
@@ -204,8 +242,7 @@ class QuotesClient:
 
   async def unsubscribe(self, instruments: List[str]) -> None:
     """
-    Unsubscribe from a list of instruments and update the subscription list.
-    Splits large unsubscription requests into batches to avoid message size issues.
+    Unsubscribe from a list of instruments.
 
     Args:
       instruments (List[str]): List of instrument identifiers.
@@ -221,7 +258,7 @@ class QuotesClient:
         unsub_list = list(to_unsubscribe)
         for batch in self._chunk_list(unsub_list, self.batch_size):
           logger.info("Unsubscribing from batch of %d instruments", len(batch))
-          message = {"action": "unsubscribe", "instruments": batch}
+          message = {"action": self.ACTION_UNSUBSCRIBE, "instruments": batch}
           await self.ws.send(json.dumps(message))
           # Small delay between batches to avoid overwhelming the server
           await asyncio.sleep(0.1)
@@ -238,6 +275,74 @@ class QuotesClient:
     """
     Close the WebSocket connection.
     """
+    self._running = False
     if self.ws:
       await self.ws.close()
-      logger.info("WebSocket connection closed.")
+      logger.info("WebSocket connection closed.")
+    
+    # Cancel the background task if it exists
+    if self._background_task and not self._background_task.done():
+      self._background_task.cancel()
+      try:
+        await self._background_task
+      except asyncio.CancelledError:
+        pass
+
+  def connect(self) -> None:
+    """
+    Connect to the websocket server in a blocking manner.
+    This method handles the event loop and will not return until stop() is called.
+    
+    Similar to KiteTicker's connect() method, this creates and runs the event loop.
+    """
+    self._running = True
+    
+    # If there's already an event loop, use it
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        # No event loop exists, create a new one
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    
+    try:
+        # Run until completion (i.e., until stop() is called)
+        loop.run_until_complete(self._connect_with_backoff())
+    finally:
+        if not loop.is_closed():
+            # Clean up pending tasks
+            pending = asyncio.all_tasks(loop)
+            for task in pending:
+                task.cancel()
+            
+            # Run until all tasks are properly canceled
+            try:
+                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
+            except Exception:
+                pass
+
+  def connect_async(self) -> None:
+    """
+    Connect to the websocket server in a non-blocking manner.
+    This method starts the connection in a background task.
+    """
+    if self._running:
+      logger.warning("Client is already running.")
+      return
+    
+    self._running = True
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+    
+    self._background_task = asyncio.create_task(self._connect_with_backoff())
+    
+  def stop(self) -> None:
+    """
+    Stop the websocket connection.
+    This is a non-blocking method that just flags the client to stop.
+    """
+    self._running = False
+    logger.info("Client stopping. Connection will close soon.")