From 646ed50f37ac6f2153efb49c8371535d0878fc09 Mon Sep 17 00:00:00 2001
From: Matthias <xmatthias@outlook.com>
Date: Thu, 15 Aug 2024 07:29:19 +0200
Subject: [PATCH] chore: improve typing for balance endpoint

---
 freqtrade/exchange/exchange.py | 4 ++--
 freqtrade/exchange/kraken.py   | 4 ++--
 freqtrade/exchange/types.py    | 8 ++++++++
 freqtrade/wallets.py           | 6 +++---
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/freqtrade/exchange/exchange.py b/freqtrade/exchange/exchange.py
index 546e6baf4..a7fe0be64 100644
--- a/freqtrade/exchange/exchange.py
+++ b/freqtrade/exchange/exchange.py
@@ -88,7 +88,7 @@ from freqtrade.exchange.exchange_utils_timeframe import (
     timeframe_to_seconds,
 )
 from freqtrade.exchange.exchange_ws import ExchangeWS
-from freqtrade.exchange.types import OHLCVResponse, OrderBook, Ticker, Tickers
+from freqtrade.exchange.types import CcxtBalances, OHLCVResponse, OrderBook, Ticker, Tickers
 from freqtrade.misc import (
     chunks,
     deep_merge_dicts,
@@ -1663,7 +1663,7 @@ class Exchange:
         return order
 
     @retrier
-    def get_balances(self) -> dict:
+    def get_balances(self) -> CcxtBalances:
         try:
             balances = self._api.fetch_balance()
             # Remove additional info from ccxt results
diff --git a/freqtrade/exchange/kraken.py b/freqtrade/exchange/kraken.py
index f0562ecaf..4b178420a 100644
--- a/freqtrade/exchange/kraken.py
+++ b/freqtrade/exchange/kraken.py
@@ -12,7 +12,7 @@ from freqtrade.enums import MarginMode, TradingMode
 from freqtrade.exceptions import DDosProtection, OperationalException, TemporaryError
 from freqtrade.exchange import Exchange
 from freqtrade.exchange.common import retrier
-from freqtrade.exchange.types import Tickers
+from freqtrade.exchange.types import CcxtBalances, Tickers
 
 
 logger = logging.getLogger(__name__)
@@ -57,7 +57,7 @@ class Kraken(Exchange):
         return super().get_tickers(symbols=symbols, cached=cached)
 
     @retrier
-    def get_balances(self) -> dict:
+    def get_balances(self) -> CcxtBalances:
         if self._config["dry_run"]:
             return {}
 
diff --git a/freqtrade/exchange/types.py b/freqtrade/exchange/types.py
index a0d315c78..564505289 100644
--- a/freqtrade/exchange/types.py
+++ b/freqtrade/exchange/types.py
@@ -25,6 +25,14 @@ class OrderBook(TypedDict):
     nonce: Optional[int]
 
 
+class CcxtBalance(TypedDict):
+    free: float
+    used: float
+    total: float
+
+
+CcxtBalances = Dict[str, CcxtBalance]
+
 Tickers = Dict[str, Ticker]
 
 # pair, timeframe, candleType, OHLCV, drop last?,
diff --git a/freqtrade/wallets.py b/freqtrade/wallets.py
index 336f24b77..f0e36b61c 100644
--- a/freqtrade/wallets.py
+++ b/freqtrade/wallets.py
@@ -139,9 +139,9 @@ class Wallets:
             if isinstance(balances[currency], dict):
                 self._wallets[currency] = Wallet(
                     currency,
-                    balances[currency].get("free"),
-                    balances[currency].get("used"),
-                    balances[currency].get("total"),
+                    balances[currency].get("free", 0),
+                    balances[currency].get("used", 0),
+                    balances[currency].get("total", 0),
                 )
         # Remove currencies no longer in get_balances output
         for currency in deepcopy(self._wallets):