Simplify using a custom network (#1327)

* Simplifies using a custom network

* update test

* update calculate_hash params type

* for chain id in account map string to hash and then to integate

* add test and getting chain_id from network in default

* fix test

* fixes after remove goerli

* feedback

* Improve logic in `Account` constructor

* Remove `chain` parameter from `Account` methods

* Add `RECOGNIZED_CHAIN_IDS` constant

* Update migration guide

---------

Co-authored-by: Dariusz Doktorski <dariusz.doktorski@swmansion.com>

Author: tkumor3

docs/migration_guide.rst

@@ -11,7 +11,17 @@ Version 0.22.0 of **starknet.py**
 -----------------------
 
 1. Support for Goerli has been removed
-2. ``StarknetChainId.SEPOLIA_TESTNET`` has been renamed to ``StarknetChainId.SEPOLIA``
+
+.. currentmodule:: starknet_py.net.models
+
+2. ``StarknetChainId.SEPOLIA_TESTNET`` has been renamed to :class:`StarknetChainId.SEPOLIA`
+
+.. currentmodule:: starknet_py.net.account.account
+
+3. Parameter ``chain`` has been removed from the methods :meth:`Account.deploy_account_v1` and :meth:`Account.deploy_account_v3`
+4. Parameter ``chain_id`` has been removed from the method :meth:`~Account.get_balance`
+
+
 
 ******************************
 0.21.0 Migration guide

starknet_py/net/account/account.py

@@ -23,7 +23,8 @@
     Tag,
 )
 from starknet_py.net.full_node_client import FullNodeClient
-from starknet_py.net.models import AddressRepresentation, StarknetChainId, parse_address
+from starknet_py.net.models import AddressRepresentation, parse_address
+from starknet_py.net.models.chains import RECOGNIZED_CHAIN_IDS, Chain, parse_chain
 from starknet_py.net.models.transaction import (
     AccountTransaction,
     DeclareV1,
@@ -72,7 +73,7 @@ def __init__(
         client: Client,
         signer: Optional[BaseSigner] = None,
         key_pair: Optional[KeyPair] = None,
-        chain: Optional[StarknetChainId] = None,
+        chain: Optional[Chain] = None,
     ):
         """
         :param address: Address of the account contract.
@@ -81,11 +82,18 @@ def __init__(
                        If none is provided, default
                        :py:class:`starknet_py.net.signer.stark_curve_signer.StarkCurveSigner` is used.
         :param key_pair: Key pair that will be used to create a default `Signer`.
-        :param chain: ChainId of the chain used to create the default signer.
+        :param chain: Chain ID associated with the account.
+            This can be supplied in multiple formats:
+
+            - an enum :py:class:`starknet_py.net.models.StarknetChainId`
+            - a string name (e.g. 'SN_SEPOLIA')
+            - a hexadecimal value (e.g. '0x1')
+            - an integer (e.g. 1)
         """
         self._address = parse_address(address)
         self._client = client
         self._cairo_version = None
+        self._chain_id = None if chain is None else parse_chain(chain)
 
         if signer is not None and key_pair is not None:
             raise ValueError("Arguments signer and key_pair are mutually exclusive.")
@@ -95,14 +103,13 @@ def __init__(
                 raise ValueError(
                     "Either a signer or a key_pair must be provided in Account constructor."
                 )
-            if chain is None:
+            if self._chain_id is None:
                 raise ValueError("One of chain or signer must be provided.")
 
             signer = StarkCurveSigner(
-                account_address=self.address, key_pair=key_pair, chain_id=chain
+                account_address=self.address, key_pair=key_pair, chain_id=self._chain_id
             )
         self.signer: BaseSigner = signer
-        self._chain_id = chain
 
     @property
     def address(self) -> int:
@@ -293,13 +300,18 @@ async def get_nonce(
     async def get_balance(
         self,
         token_address: Optional[AddressRepresentation] = None,
-        chain_id: Optional[StarknetChainId] = None,
         *,
         block_hash: Optional[Union[Hash, Tag]] = None,
         block_number: Optional[Union[int, Tag]] = None,
     ) -> int:
         if token_address is None:
-            token_address = self._default_token_address_for_chain(chain_id)
+            chain_id = await self._get_chain_id()
+            if chain_id in RECOGNIZED_CHAIN_IDS:
+                token_address = FEE_CONTRACT_ADDRESS
+            else:
+                raise ValueError(
+                    "Argument token_address must be specified when using a custom network."
+                )
 
         low, high = await self._client.call_contract(
             Call(
@@ -596,7 +608,6 @@ async def deploy_account_v1(
         salt: int,
         key_pair: KeyPair,
         client: Client,
-        chain: StarknetChainId,
         constructor_calldata: Optional[List[int]] = None,
         nonce: int = 0,
         max_fee: Optional[int] = None,
@@ -618,7 +629,6 @@ async def deploy_account_v1(
         :param salt: Salt used to calculate the address.
         :param key_pair: KeyPair used to calculate address and sign deploy account transaction.
         :param client: Client instance used for deployment.
-        :param chain: Id of the Starknet chain used.
         :param constructor_calldata: Optional calldata to account contract constructor. If ``None`` is passed,
             ``[key_pair.public_key]`` will be used as calldata.
         :param nonce: Nonce of the transaction.
@@ -631,6 +641,8 @@ async def deploy_account_v1(
             else [key_pair.public_key]
         )
 
+        chain = await client.get_chain_id()
+
         account = _prepare_account_to_deploy(
             address=address,
             class_hash=class_hash,
@@ -650,7 +662,7 @@ async def deploy_account_v1(
             auto_estimate=auto_estimate,
         )
 
-        if chain in StarknetChainId:
+        if parse_chain(chain) in RECOGNIZED_CHAIN_IDS:
             balance = await account.get_balance()
             if balance < deploy_account_tx.max_fee:
                 raise ValueError(
@@ -671,7 +683,6 @@ async def deploy_account_v3(
         salt: int,
         key_pair: KeyPair,
         client: Client,
-        chain: StarknetChainId,
         constructor_calldata: Optional[List[int]] = None,
         nonce: int = 0,
         l1_resource_bounds: Optional[ResourceBounds] = None,
@@ -690,7 +701,6 @@ async def deploy_account_v3(
         :param salt: Salt used to calculate the address.
         :param key_pair: KeyPair used to calculate address and sign deploy account transaction.
         :param client: Client instance used for deployment.
-        :param chain: Id of the Starknet chain used.
         :param constructor_calldata: Optional calldata to account contract constructor. If ``None`` is passed,
             ``[key_pair.public_key]`` will be used as calldata.
         :param nonce: Nonce of the transaction.
@@ -704,6 +714,8 @@ async def deploy_account_v3(
             else [key_pair.public_key]
         )
 
+        chain = await client.get_chain_id()
+
         account = _prepare_account_to_deploy(
             address=address,
             class_hash=class_hash,
@@ -729,19 +741,12 @@ async def deploy_account_v3(
             hash=result.transaction_hash, account=account, _client=account.client
         )
 
-    def _default_token_address_for_chain(
-        self, chain_id: Optional[StarknetChainId] = None
-    ) -> str:
-        if (chain_id or self._chain_id) not in [
-            StarknetChainId.SEPOLIA,
-            StarknetChainId.SEPOLIA_INTEGRATION,
-            StarknetChainId.MAINNET,
-        ]:
-            raise ValueError(
-                "Argument token_address must be specified when using a custom network."
-            )
+    async def _get_chain_id(self) -> int:
+        if self._chain_id is None:
+            chain = await self._client.get_chain_id()
+            self._chain_id = parse_chain(chain)
 
-        return FEE_CONTRACT_ADDRESS
+        return self._chain_id
 
 
 def _prepare_account_to_deploy(
@@ -750,7 +755,7 @@ def _prepare_account_to_deploy(
     salt: int,
     key_pair: KeyPair,
     client: Client,
-    chain: StarknetChainId,
+    chain: Chain,
     calldata: List[int],
 ) -> Account:
     # pylint: disable=too-many-arguments

starknet_py/net/account/account_test.py

@@ -88,27 +88,42 @@ def test_create_account():
     assert account.signer.public_key == key_pair.public_key
 
 
-def test_create_account_from_signer():
+@pytest.mark.parametrize(
+    "chain",
+    [
+        StarknetChainId.SEPOLIA,
+        "SN_SEPOLIA",
+        "0x534e5f5345504f4c4941",
+        393402133025997798000961,
+    ],
+)
+def test_create_account_parses_chain(chain):
+    key_pair = KeyPair.from_private_key(0x111)
+    account = Account(
+        address=0x1,
+        client=FullNodeClient(node_url=""),
+        key_pair=key_pair,
+        chain=chain,
+    )
+
+    assert account.address == 0x1
+    assert account.signer.public_key == key_pair.public_key
+    assert isinstance(account.signer, StarkCurveSigner)
+    assert account.signer.chain_id == 0x534E5F5345504F4C4941
+
+
+def test_create_account_from_signer(client):
     signer = StarkCurveSigner(
         account_address=0x1,
         key_pair=KeyPair.from_private_key(0x111),
         chain_id=StarknetChainId.MAINNET,
     )
-    account = Account(address=0x1, client=FullNodeClient(node_url=""), signer=signer)
+    account = Account(address=0x1, client=client, signer=signer)
 
     assert account.address == 0x1
     assert account.signer == signer
 
 
-def test_create_account_raises_on_no_chain_and_signer():
-    with pytest.raises(ValueError, match="One of chain or signer must be provided"):
-        Account(
-            address=0x1,
-            client=FullNodeClient(node_url=""),
-            key_pair=KeyPair.from_private_key(0x111),
-        )
-
-
 def test_create_account_raises_on_no_keypair_and_signer():
     with pytest.raises(
         ValueError,

starknet_py/net/account/base_account.py

@@ -10,7 +10,7 @@
     SentTransactionResponse,
     Tag,
 )
-from starknet_py.net.models import AddressRepresentation, StarknetChainId
+from starknet_py.net.models import AddressRepresentation
 from starknet_py.net.models.transaction import (
     AccountTransaction,
     DeclareV1,
@@ -93,18 +93,15 @@ async def get_nonce(
     async def get_balance(
         self,
         token_address: Optional[AddressRepresentation] = None,
-        chain_id: Optional[StarknetChainId] = None,
         *,
         block_hash: Optional[Union[Hash, Tag]] = None,
         block_number: Optional[Union[int, Tag]] = None,
     ) -> int:
         """
-        Checks account's balance of specified token.
+        Checks account's balance of the specified token.
+        By default, it uses the L2 ETH address for mainnet and sepolia networks.
 
         :param token_address: Address of the ERC20 contract.
-        :param chain_id: Identifier of the Starknet chain used.
-            If token_address is not specified it will be used to determine network's payment token address.
-            If token_address is provided, chain_id will be ignored.
         :param block_hash: Block's hash or literals `"pending"` or `"latest"`
         :param block_number: Block's number or literals `"pending"` or `"latest"`
         :return: Token balance.

starknet_py/net/client.py

@@ -301,3 +301,7 @@ async def get_contract_nonce(
         :param block_number: Block's number or literals `"pending"` or `"latest"`
         :return: The last nonce used for the given contract
         """
+
+    @abstractmethod
+    async def get_chain_id(self) -> str:
+        """Return the currently configured Starknet chain id"""

starknet_py/net/full_node_client.py

@@ -455,7 +455,6 @@ async def get_block_hash_and_number(self) -> BlockHashAndNumber:
         return cast(BlockHashAndNumber, BlockHashAndNumberSchema().load(res))
 
     async def get_chain_id(self) -> str:
-        """Return the currently configured Starknet chain id"""
         return await self._client.call(method_name="chainId", params={})
 
     async def get_syncing_status(self) -> Union[bool, SyncStatus]:

starknet_py/net/models/chains.py

@@ -1,5 +1,5 @@
 from enum import IntEnum
-from typing import Optional
+from typing import Optional, Union
 
 from starknet_py.common import int_from_bytes
 from starknet_py.net.networks import MAINNET, SEPOLIA, SEPOLIA_INTEGRATION, Network
@@ -15,6 +15,13 @@ class StarknetChainId(IntEnum):
     SEPOLIA_INTEGRATION = int_from_bytes(b"SN_INTEGRATION_SEPOLIA")
 
 
+RECOGNIZED_CHAIN_IDS = [
+    StarknetChainId.MAINNET,
+    StarknetChainId.SEPOLIA,
+    StarknetChainId.SEPOLIA_INTEGRATION,
+]
+
+
 def chain_from_network(
     net: Network, chain: Optional[StarknetChainId] = None
 ) -> StarknetChainId:
@@ -31,3 +38,17 @@ def chain_from_network(
         raise ValueError("Chain is required when not using predefined networks.")
 
     return chain
+
+
+ChainId = Union[StarknetChainId, int]
+Chain = Union[str, ChainId]
+
+
+def parse_chain(chain: Chain) -> ChainId:
+    if isinstance(chain, str):
+        try:
+            return int(chain, 16)
+        except ValueError:
+            return int_from_bytes(chain.encode())
+    else:
+        return chain