Make native balance optional for ERC20 get_balances

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 [
  {

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

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