feat(wallet): Add BaseWallet.get_balance_per_address() method

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)

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
         """

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