add balance subscription in account websocket client

moves key building of subscription to each class

Author: Insmael

cryptomarket/websockets/account_client.py

@@ -12,8 +12,30 @@ class AccountClient(ClientAuth):
     :param on_error: function called on a websocket error, and called in an authenticated error. it takes one parameter, the error.
     :param on_close: function called on the closing event of the websocket. no parameters
     """
-    def __init__(self, api_key:str, api_secret:str, on_connect:callable=None, on_error:callable=None, on_close:callable=None):
-        super(AccountClient, self).__init__("wss://api.exchange.cryptomkt.com/api/2/ws/account", api_key, api_secret, on_connect=on_connect, on_error=on_error, on_close=on_close)
+    def __init__(
+        self, 
+        api_key:str, 
+        api_secret:str, 
+        on_connect:callable=None, 
+        on_error:callable=None, 
+        on_close:callable=None
+    ):
+        super(AccountClient, self).__init__(
+            "wss://api.exchange.cryptomkt.com/api/2/ws/account", 
+            api_key, 
+            api_secret, 
+            subscription_keys={
+                "unsubscribeBalance":"balance", 
+                "subscribeBalance":"balance",
+                "balance":"balance",
+                "subscribeTransactions":"transaction",
+                "subscribeTransactions":"transaction",
+                "updateTransaction":"transaction"
+            },
+            on_connect=on_connect, 
+            on_error=on_error, 
+            on_close=on_close
+        )
 
     def get_account_balance(self, callback: callable):
         """Get the user account balance.
@@ -46,15 +68,17 @@ def get_account_balance(self, callback: callable):
         """
         self.send_by_id(method='getBalance', callback=callback)
 
-    def find_transactions(self,
+    def find_transactions(
+        self,
         callback: callable,
         currency: str = None,
         sort: str = None,
         since: str = None,
         till: str = None,
         limit: int = None,
         offset: int = None,
-        show_senders: bool = None):
+        show_senders: bool = None
+    ):
         """Get a list of transactions of the account. Accepts only filtering by Datetime
         
         https://api.exchange.cryptomkt.com/#find-transactions
@@ -87,15 +111,17 @@ def find_transactions(self,
         params = args.DictBuilder().currency(currency).sort(sort).since(since).till(till).limit(limit).offset(offset).show_senders(show_senders).build()
         self.send_by_id("findTransactions", callback, params)
 
-    def load_transactions(self,
+    def load_transactions(
+        self,
         callback: callable,
         currency: str = None,
         sort: str = None,
         since: str = None,
         till: str = None,
         limit: int = None,
         offset: int = None,
-        show_senders: bool = None):
+        show_senders: bool = None
+    ):
         """Get a list of transactions of the account. Accepts only filtering by Index.
 
         https://api.exchange.cryptomkt.com/#load-transactions
@@ -175,3 +201,42 @@ def unsubscribe_to_transactions(self, callback: callable=None):
         :returns: The operation result as result argument for the callback. True if success.
         """
         self.send_unsubscription(method='unsubscribeTransactions', callback=callback)
+
+    def subscribe_to_balance(self, callback: callable, result_callback: callable=None):
+        """Subscribe to the balance of the account.
+
+        This subscription aims to provide an easy way to be informed of the current balance state. 
+        If the state has been changed or potentially changed the "balance" event will come with the actual state. 
+        Please be aware that only non-zero values are present.
+
+        https://api.exchange.cryptomkt.com/#subscription-to-the-balance
+
+        :param callback: A callable to call with each update of the result data. It takes one argument, the balance feed.
+        :param result_callback: A callable to call with the subscription result. It takes two arguments, err and result. err is None for successful calls, result is None for calls with error: callback(err, result).
+
+        :returns: A transaction of the account as feed for the callback.
+
+        .. code-block:: python
+        {
+            "id": "76b70d1c-3dd7-423e-976e-902e516aae0e",
+            "index": 7173627250,
+            "type": "bankToExchange",
+            "status": "success",
+            "currency": "BTG",
+            "amount": "0.00001000",
+            "createdAt": "2021-01-31T08:19:33.892Z",
+            "updatedAt": "2021-01-31T08:19:33.967Z"
+        }
+        """
+        self.send_subscription(method='subscribeBalance', callback=callback, params={}, result_callback=result_callback)
+
+    def unsubscribe_to_balance(self, callback: callable=None):
+        """unsubscribe to the balance feed.
+
+        https://api.exchange.cryptomkt.com/#subscription-to-the-balance
+
+        :param callback: Optional. A callable to call with the result data. It takes two arguments, err and result. err is None for successful calls, result is None for calls with error: callback(err, result).
+
+        :returns: The operation result as result argument for the callback. True if success.
+        """
+        self.send_unsubscription(method='unsubscribeBalance', callback=callback)

cryptomarket/websockets/client_auth.py

@@ -8,8 +8,23 @@
 
 
 class ClientAuth(ClientBase):
-    def __init__(self, uri:str, api_key:str, api_secret:str, on_connect=None, on_error=None, on_close=None):
-        super(ClientAuth, self).__init__(uri, on_connect=on_connect, on_error=on_error, on_close=on_close)
+    def __init__(
+        self, 
+        uri:str, 
+        api_key:str, 
+        api_secret:str, 
+        subscription_keys:Dict[str, str]={}, 
+        on_connect=None, 
+        on_error=None, 
+        on_close=None
+    ):
+        super(ClientAuth, self).__init__(
+            uri, 
+            subscription_keys=subscription_keys, 
+            on_connect=on_connect, 
+            on_error=on_error, 
+            on_close=on_close
+        )
         self.api_key = api_key
         self.api_secret = api_secret
         self.authed = False

cryptomarket/websockets/client_base.py

@@ -3,12 +3,21 @@
 from cryptomarket.exceptions import CryptomarketAPIException
 from cryptomarket.websockets.callback_cache import CallbackCache
 from cryptomarket.websockets.manager import WebsocketManager
+from types import Dict
 
 
 class ClientBase:
-    def __init__(self, uri:str, on_connect=None, on_error=None, on_close=None):
+    def __init__(
+        self, 
+        uri:str, 
+        subscription_keys:Dict[str,str]={}, 
+        on_connect=None, 
+        on_error=None, 
+        on_close=None
+    ):
         self.ws_manager = WebsocketManager(self, uri)    
-        self.callback_cache = CallbackCache()    
+        self.callback_cache = CallbackCache()
+        self.subscription_keys = subscription_keys
 
         if on_connect is not None: self.on_connect = on_connect
         else: self.on_connect = lambda: None
@@ -90,4 +99,6 @@ def handle_response(self, response):
                 callback(None, result)
 
     def build_key(self, method, params):
-        return "subscription"
+        if not method in self.subscription_keys: 
+            return "subscription"
+        return self.subscription_map[method]

cryptomarket/websockets/methods.py

@@ -2,8 +2,6 @@
 
 import cryptomarket.args as args
 from cryptomarket.websockets.client_base import ClientBase
-from cryptomarket.websockets.methods import (candles_feed, method_key,
-                                             orderbook_feed, trades_feed)
 from cryptomarket.websockets.orderbook_cache import OrderbookCache
 
 
@@ -13,9 +11,43 @@ class PublicClient(ClientBase):
     :param callback: A callable to call with the client once the connection is established. if an error ocurrs is return as the fist parameter of the callback: callback(err, client)
     """
     def __init__(self, on_connect=None, on_error=None, on_close=None):
-        super(PublicClient, self).__init__("wss://api.exchange.cryptomkt.com/api/2/ws/public", on_connect=on_connect, on_error=on_error, on_close=on_close)
+        super(PublicClient, self).__init__(
+            "wss://api.exchange.cryptomkt.com/api/2/ws/public", 
+            subscription_keys={
+                # tickers
+                'subscribeTicker':'tickers',
+                'unsubscribeTicker':'tickers',
+                'ticker':'tickers',
+                # orderbooks
+                'subscribeOrderbook':'orderbooks',
+                'unsubscribeOrderbook':'orderbooks',
+                'snapshotOrderbook':'orderbooks',
+                'updateOrderbook':'orderbooks',
+                # trades
+                'subscribeTrades':'trades',
+                'unsubscribeTrades':'trades',
+                'snapshotTrades':'trades',
+                'updateTrades':'trades',
+                # candles
+                'subscribeCandles':'candles',
+                'unsubscribeCandles':'candles',
+                'snapshotCandles':'candles',
+                'updateCandles':'candles',
+            },
+            on_connect=on_connect, 
+            on_error=on_error, 
+            on_close=on_close
+        )
         self.ob_cache = OrderbookCache()
+    
+    def orderbook_feed(self, method):
+        return method == 'snapshotOrderbook' or method == 'updateOrderbook'
+
+    def candles_feed(self, method):
+        return method == 'snapshotCandles' or method == 'updateCandles'
 
+    def trades_feed(self, method):
+        return method == 'snapshotTrades' or method == 'updateTrades'
 
     def handle_notification(self, message):
         params = message['params']
@@ -24,15 +56,15 @@ def handle_notification(self, message):
         callback = self.callback_cache.get_subscription_callback(key)
         if callback is None: return
         subscription_data = None
-        if orderbook_feed(method):
+        if self.orderbook_feed(method):
             self.ob_cache.update(method, key, params)
             if self.ob_cache.orderbook_broken(key):
                 self.send_by_id('subscribeOrderbook', None, {'symbol':symbol})
                 self.ob_cache.wait_orderbook(key)
             if self.ob_cache.orderbook_wating(key): return
             subscription_data = self.ob_cache.get_ob(key)
 
-        elif trades_feed(method) or candles_feed(method):
+        elif self.trades_feed(method) or self.candles_feed(method):
             subscription_data = params['data']
 
         else:
@@ -41,7 +73,7 @@ def handle_notification(self, message):
         callback(subscription_data)
 
     def build_key(self, method, params):
-        m_key = method_key(method)
+        m_key = self.subscription_keys[method]
         symbol = params['symbol'] if 'symbol' in params else ''
         period = params['period'] if 'period' in params else ''
         key = m_key + ':' + symbol + ':' + period

cryptomarket/websockets/public_client.py

@@ -3,7 +3,6 @@
 import cryptomarket.args as args
 from cryptomarket.websockets.client_auth import ClientAuth
 
-
 class TradingClient(ClientAuth):
     """TradingClient connects via websocket to cryptomarket to enable the user to manage orders. uses SHA256 as auth method and authenticates automatically.
 
@@ -13,22 +12,43 @@ class TradingClient(ClientAuth):
     :param on_error: function called on a websocket error, and called in an authenticated error. it takes one parameter, the error.
     :param on_close: function called on the closing event of the websocket. no parameters
     """
-    def __init__(self, api_key:str, api_secret:str, on_connect:callable=None, on_error:callable=None, on_close:callable=None):
-        super(TradingClient, self).__init__("wss://api.exchange.cryptomkt.com/api/2/ws/trading", api_key, api_secret, on_connect=on_connect, on_error=on_error, on_close=on_close)
-
-    def create_order(self,
-    client_order_id: str, 
-    symbol: str, 
-    side: str, 
-    quantity: float,
-    order_type: str = None,
-    price: float = None,
-    stop_price: float = None,
-    time_in_force: str = None,
-    expire_time: str = None,
-    strict_validate: bool = False,
-    post_only: bool = False,
-    callback: callable = None):
+    def __init__(
+        self, 
+        api_key:str, 
+        api_secret:str, 
+        on_connect:callable=None, 
+        on_error:callable=None, 
+        on_close:callable=None
+    ):
+        super(TradingClient, self).__init__(
+            "wss://api.exchange.cryptomkt.com/api/2/ws/trading", 
+            api_key, 
+            api_secret, 
+            subscription_keys={
+                "subscribeReports":"reports",
+                'activeOrders':'reports',
+                'report':'reports',
+            },
+            on_connect=on_connect, 
+            on_error=on_error, 
+            on_close=on_close
+        )
+
+    def create_order(
+        self,
+        client_order_id: str, 
+        symbol: str, 
+        side: str, 
+        quantity: float,
+        order_type: str = None,
+        price: float = None,
+        stop_price: float = None,
+        time_in_force: str = None,
+        expire_time: str = None,
+        strict_validate: bool = False,
+        post_only: bool = False,
+        callback: callable = None
+    ):
         """Creates a new order
 
         Requires authentication.
@@ -73,9 +93,11 @@ def create_order(self,
         params = builder.strict_validate(strict_validate).post_only(post_only).build()
         self.send_by_id(method='newOrder', callback=callback, params=params)
 
-    def cancel_order(self, 
-    client_order_id: str, 
-    callback: callable= None):
+    def cancel_order(
+        self, 
+        client_order_id: str, 
+        callback: callable= None
+    ):
         """Cancel the order with client_order_id.
 
         Requires authentication.