diff --git a/gnosis/eth/ethereum_client.py b/gnosis/eth/ethereum_client.py index 377650d4c..eee23ba16 100644 --- a/gnosis/eth/ethereum_client.py +++ b/gnosis/eth/ethereum_client.py @@ -517,13 +517,17 @@ def get_balance( ) def get_balances( - self, address: ChecksumAddress, token_addresses: Sequence[ChecksumAddress] + self, + address: ChecksumAddress, + token_addresses: Sequence[ChecksumAddress], + include_native_balance: bool = True, ) -> List[BalanceDict]: """ Get balances for Ether and tokens for an `address` :param address: Owner address checksummed :param token_addresses: token addresses to check + :param include_native_balance: if `True` returns also the native token balance :return: ``List[BalanceDict]`` """ @@ -543,6 +547,9 @@ def get_balances( for token_address, balance in zip(token_addresses, balances) ] + if not include_native_balance: + return return_balances + # Add ether balance response return [ { diff --git a/gnosis/eth/tests/test_ethereum_client.py b/gnosis/eth/tests/test_ethereum_client.py index d6447c08d..1a2755212 100644 --- a/gnosis/eth/tests/test_ethereum_client.py +++ b/gnosis/eth/tests/test_ethereum_client.py @@ -207,6 +207,18 @@ def test_get_balances(self): ], ) + self.assertCountEqual( + self.ethereum_client.erc20.get_balances( + account_address, + [erc20.address, erc20_2.address], + include_native_balance=False, + ), + [ + {"token_address": erc20.address, "balance": tokens_value}, + {"token_address": erc20_2.address, "balance": tokens_value_2}, + ], + ) + with mock.patch.object( EthereumClient, "batch_call_same_function", diff --git a/gnosis/eth/typing.py b/gnosis/eth/typing.py index 2196e4360..c79250bc8 100644 --- a/gnosis/eth/typing.py +++ b/gnosis/eth/typing.py @@ -1,6 +1,6 @@ from typing import Any, Dict, Optional, TypedDict, Union -from eth_typing import Hash32, HexStr +from eth_typing import ChecksumAddress, Hash32, HexStr from hexbytes import HexBytes from web3.types import LogReceipt @@ -9,7 +9,7 @@ class BalanceDict(TypedDict): - token_address: Optional[str] + token_address: Optional[ChecksumAddress] balance: int