Add wallet manager to track wallet updates/snapshots

bfxapi/models/Wallet.py

@@ -0,0 +1,19 @@
+
+class Wallet:
+
+  def __init__(self, wType, currency, balance, unsettled_interest):
+    self.type = wType
+    self.currency = currency
+    self.balance = balance
+    self.unsettled_interest = unsettled_interest
+    self.key = "{}_{}".format(wType, currency)
+
+  def set_balance(self, data):
+    self.balance = data
+
+  def set_unsettled_interest(self, data):
+    self.unsettled_interest = data
+
+  def __str__(self):
+    return "Wallet <'{}_{}' balance='{}' unsettled='{}'>".format(
+      self.type, self.currency, self.balance, self.unsettled_interest)

bfxapi/models/__init__.py

@@ -4,3 +4,4 @@ from .Order import *
 from .Trade import *
 from .OrderBook import *
 from .Subscription import *
+from .Wallet import *

bfxapi/websockets/BfxWebsocket.py

@@ -7,6 +7,7 @@ import random
 
 from .GenericWebsocket import GenericWebsocket, AuthError
 from .SubscriptionManager import SubscriptionManager
+from .WalletManager import WalletManager
 from .OrderManager import OrderManager
 from ..models import Order, Trade, OrderBook
 
@@ -144,6 +145,7 @@ class BfxWebsocket(GenericWebsocket):
     super(BfxWebsocket, self).__init__(host, *args, **kwargs)
     self.subscriptionManager = SubscriptionManager(self)
     self.orderManager = OrderManager(self)
+    self.wallets = WalletManager()
 
     self._WS_DATA_HANDLERS = {
       'tu': self._trade_update_handler,
@@ -249,9 +251,9 @@ class BfxWebsocket(GenericWebsocket):
   
   async def _wallet_update_handler(self, data):
     # [0,"wu",["exchange","USD",89134.66933283,0]]
-    wu = data[2]
+    uw = self.wallets._update_from_event(data)
-    self._emit('wallet_update', data)
+    self._emit('wallet_update', uw)
-    self.logger.info("Wallet update: {}({}) = {}".format(wu[1], wu[0], wu[2]))
+    self.logger.info("Wallet update: {}".format(uw))
 
   async def _heart_beat_handler(self, data):
     self.logger.debug("Heartbeat - {}".format(self.host))
@@ -293,12 +295,11 @@ class BfxWebsocket(GenericWebsocket):
     await self.orderManager.confirm_order_new(data)
 
   async def _order_snapshot_handler(self, data):
-    self._emit('order_snapshot', data)
+    await self.orderManager.build_from_order_snapshot(data)
-    self.logger.info("Position snapshot: {}".format(data))
   
   async def _wallet_snapshot_handler(self, data):
-    self._emit('wallet_snapshot', data[2])
+    wallets = self.wallets._update_from_snapshot(data)
-    self.logger.info("Wallet snapshot: {}".format(data))
+    self._emit('wallet_snapshot', wallets)
   
   async def _position_snapshot_handler(self, data):
     self._emit('position_snapshot', data)
@@ -447,8 +448,11 @@ class BfxWebsocket(GenericWebsocket):
   async def update_order(self, *args, **kwargs):
     return await self.orderManager.update_order(*args, **kwargs)
 
-  async def cancel_order(self, *args, **kwargs):
+  async def close_order(self, *args, **kwargs):
-    return await self.orderManager.cancel_order(*args, **kwargs)
+    return await self.orderManager.close_order(*args, **kwargs)
+
+  async def close_all_orders(self, *args, **kwargs):
+    return await self.orderManager.close_all_orders(*args, **kwargs)
   
-  async def cancel_order_multi(self, *args, **kwargs):
+  async def close_order_multi(self, *args, **kwargs):
-    return await self.cancel_order_multi(*args, **kwargs)
+    return await self.close_order_multi(*args, **kwargs)

bfxapi/websockets/WalletManager.py

@@ -0,0 +1,27 @@
+
+from ..models import Wallet
+
+class WalletManager:
+
+  def __init__(self):
+    self.wallets = {}
+
+  def _update_from_snapshot(self, raw_ws_data):
+    # [0, 'ws', [['exchange', 'BTC', 41.25809589, 0, None], ['exchange', 'USD', 62761.86070104, 0, None]]]
+    wData = raw_ws_data[2]
+    self.wallets = {}
+    for wallet in wData:
+      new_wallet = Wallet(wallet[0], wallet[1], wallet[2], wallet[3])
+      self.wallets[new_wallet.key] = new_wallet
+    return self.get_wallets()
+
+  def _update_from_event(self, raw_ws_data):
+    # [0,"wu",["exchange","USD",62761.86070104,0,61618.66070104]]
+    wallet = raw_ws_data[2]
+    new_wallet = Wallet(wallet[0], wallet[1], wallet[2], wallet[3])
+    self.wallets[new_wallet.key] = new_wallet
+    return new_wallet
+
+  def get_wallets(self):
+    return list(self.wallets.values())
+