From 32da24dd339c1440aca24ee3018aa23abc866d3e Mon Sep 17 00:00:00 2001
From: Marcelo Salhab Brogliato <msbrogli@gmail.com>
Date: Thu, 21 Sep 2023 17:53:36 -0500
Subject: [PATCH] feat(wallet): Add BaseWallet.get_balance_per_address() method

---
 hathor/types.py                     | 1 +
 hathor/wallet/base_wallet.py        | 8 ++++++++
 tests/wallet/test_balance_update.py | 3 +++
 3 files changed, 12 insertions(+)

diff --git a/hathor/types.py b/hathor/types.py
index 07c42194a..f035a8c80 100644
--- a/hathor/types.py
+++ b/hathor/types.py
@@ -17,6 +17,7 @@
 
 VertexId = bytes        # NewType('TxId', bytes)
 Address = bytes         # NewType('Address', bytes)
+AddressB58 = str
 TxOutputScript = bytes  # NewType('TxOutputScript', bytes)
 Timestamp = int         # NewType('Timestamp', int)
 TokenUid = VertexId     # NewType('TokenUid', VertexId)
diff --git a/hathor/wallet/base_wallet.py b/hathor/wallet/base_wallet.py
index 63c7287c9..913c13bc7 100644
--- a/hathor/wallet/base_wallet.py
+++ b/hathor/wallet/base_wallet.py
@@ -32,6 +32,7 @@
 from hathor.transaction.scripts import P2PKH, create_output_script, parse_address_script
 from hathor.transaction.storage import TransactionStorage
 from hathor.transaction.transaction import Transaction
+from hathor.types import AddressB58, Amount, TokenUid
 from hathor.util import Reactor
 from hathor.wallet.exceptions import InputDuplicated, InsufficientFunds, PrivateKeyNotFound
 
@@ -135,6 +136,13 @@ def __init__(self, directory: str = './', pubsub: Optional[PubSubManager] = None
     def _manually_initialize(self) -> None:
         pass
 
+    def get_balance_per_address(self, token_uid: TokenUid) -> dict[AddressB58, Amount]:
+        """Return balance per address for a given token. This method ignores locks."""
+        balances: defaultdict[AddressB58, Amount] = defaultdict(Amount)
+        for utxo_info, unspent_tx in self.unspent_txs[token_uid].items():
+            balances[unspent_tx.address] += unspent_tx.value
+        return dict(balances)
+
     def start(self) -> None:
         """ Start the pubsub subscription if wallet has a pubsub
         """
diff --git a/tests/wallet/test_balance_update.py b/tests/wallet/test_balance_update.py
index ab2260866..15aad3e6b 100644
--- a/tests/wallet/test_balance_update.py
+++ b/tests/wallet/test_balance_update.py
@@ -436,6 +436,9 @@ def test_tokens_balance(self):
         # hathor balance remains the same
         self.assertEqual(self.manager.wallet.balance[settings.HATHOR_TOKEN_UID], hathor_balance)
 
+        balances_per_address = self.manager.wallet.get_balance_per_address(settings.HATHOR_TOKEN_UID)
+        self.assertEqual(hathor_balance.available, sum(x for x in balances_per_address.values()))
+
 
 class SyncV1HathorSyncMethodsTestCase(unittest.SyncV1Params, BaseHathorSyncMethodsTestCase):
     __test__ = True