bfxbot/bfxwrapper.py

@@ -1,4 +1,5 @@
 from enum import Enum
+from typing import List
 
 from bfxapi.rest.bfx_rest import BfxRest
 from retrying_async import retry
@@ -6,6 +7,29 @@ from retrying_async import retry
 from bfxbot.currency import Symbol
 
 
+class BalanceKind(Enum):
+    EXCHANGE = "exchange",
+    MARGIN = "margin"
+
+    @staticmethod
+    def from_str(string: str):
+        string = string.lower()
+
+        if "margin" in string:
+            return BalanceKind.MARGIN
+        if "exchange" in string:
+            return BalanceKind.EXCHANGE
+
+        return None
+
+
+class Balance:
+    def __init__(self, currency: str, amount: int, kind: BalanceKind):
+        self.currency: str = currency
+        self.amount: int = amount
+        self.kind: BalanceKind = kind
+
+
 class Direction(Enum):
     UP = 1,
     DOWN = -1
@@ -50,8 +74,8 @@ class BfxWrapper(BfxRest):
         return await super().get_trades(symbol, start, end)
 
     @retry()
-    async def post(self, api_path: str, data: dict):
-        return await super().post(api_path, data)
+    async def post(self, endpoint: str, data: dict = {}, params=""):
+        return await super().post(endpoint, data, params)
 
     ################################
     # NEW METHODS
@@ -83,6 +107,20 @@ class BfxWrapper(BfxRest):
 
         return balance
 
+    async def get_balances(self) -> List[Balance]:
+        balances = []
+        wallets = await self.get_wallets()
+
+        for w in wallets:
+            kind = BalanceKind.from_str(w.type)
+
+            if not kind:
+                continue
+
+            balances.append(Balance(w.currency, w.balance, kind))
+
+        return balances
+
     async def maximum_order_amount(self, symbol: Symbol, direction: Direction,
                                    order_type: OrderType = OrderType.EXCHANGE,
                                    rate: int = 1):