refactor: simplify the ChainType

valory-xyz · Aug 2, 2024 · e6f524a · e6f524a
1 parent d503451
commit e6f524a
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 41 deletions.
diff --git a/operate/cli.py b/operate/cli.py
@@ -349,7 +349,7 @@ async def _get_wallets(request: Request) -> t.List[t.Dict]:
     async def _get_wallet_by_chain(request: Request) -> t.List[t.Dict]:
         """Create wallet safe"""
         ledger_type = get_ledger_type_from_chain_type(
-            chain=ChainType.from_string(request.path_params["chain"])
+            chain=ChainType[request.path_params["chain"].upper()]
         )
         manager = operate.wallet_manager
         if not manager.exists(ledger_type=ledger_type):
@@ -405,7 +405,7 @@ async def _get_safes(request: Request) -> t.List[t.Dict]:
     async def _get_safe(request: Request) -> t.List[t.Dict]:
         """Create wallet safe"""
         ledger_type = get_ledger_type_from_chain_type(
-            chain=ChainType.from_string(request.path_params["chain"])
+            chain=ChainType[request.path_params["chain"].upper()]
         )
         manager = operate.wallet_manager
         if not manager.exists(ledger_type=ledger_type):

diff --git a/operate/services/service.py b/operate/services/service.py
@@ -244,7 +244,7 @@ def ledger_config(self) -> "LedgerConfig":
                 (_, config), *_ = override["config"]["ledger_apis"].items()
                 return LedgerConfig(
                     rpc=config["address"],
-                    chain=ChainType.from_id(cid=config["chain_id"]),
+                    chain=ChainType(config["chain_id"]),
                     type=LedgerType.ETHEREUM,
                 )
         raise ValueError("No ledger config found.")

diff --git a/operate/types.py b/operate/types.py
@@ -35,23 +35,6 @@
     "stop": 3,
 }
 
-
-_CHAIN_NAME_TO_ENUM = {
-    "ethereum": 0,
-    "goerli": 1,
-    "gnosis": 2,
-    "solana": 3,
-}
-
-_CHAIN_ID_TO_CHAIN_NAME = {
-    1: "ethereum",
-    5: "goerli",
-    100: "gnosis",
-    1399811149: "solana",
-}
-
-_CHAIN_NAME_TO_ID = {val: key for key, val in _CHAIN_ID_TO_CHAIN_NAME.items()}
-
 _LEDGER_TYPE_TO_ENUM = {
     "ethereum": 0,
     "solana": 1,
@@ -83,25 +66,21 @@ def key_file(self) -> str:
 class ChainType(enum.IntEnum):
     """Chain type enum."""
 
-    ETHEREUM = 0
-    GOERLI = 1
-    GNOSIS = 2
-    SOLANA = 3
-
-    @property
-    def id(self) -> int:
-        """Returns chain id."""
-        return _CHAIN_NAME_TO_ID[self.name.lower()]
-
-    @classmethod
-    def from_string(cls, chain: str) -> "ChainType":
-        """Load from string."""
-        return cls(_CHAIN_NAME_TO_ENUM[chain.lower()])
-
-    @classmethod
-    def from_id(cls, cid: int) -> "ChainType":
-        """Load from chain ID."""
-        return cls(_CHAIN_NAME_TO_ENUM[_CHAIN_ID_TO_CHAIN_NAME[cid]])
+    ETHEREUM = 1
+    GOERLI = 5
+    GNOSIS = 100
+    SOLANA = 1399811149
+    POLYGON = 137
+    POLYGON_MUMBAI = 80001
+    CHIADO = 10200
+    ARBITRUM_ONE = 42161
+    ARBITRUM_SEPOLIA = 421614
+    OPTIMISTIC = 10
+    OPTIMISTIC_SEPOLIA = 11155420
+    BASE = 8453
+    BASE_SEPOLIA = 84532
+    CELO = 42220
+    CELO_ALFAJORES = 42220
 
 
 class Action(enum.IntEnum):

diff --git a/operate/wallet/master.py b/operate/wallet/master.py
@@ -90,7 +90,7 @@ def ledger_api(
         return make_ledger_api(
             self.ledger_type.name.lower(),
             address=(rpc or get_default_rpc(chain=chain_type)),
-            chain_id=chain_type.id,
+            chain_id=chain_type.value,
         )
 
     def transfer(
@@ -184,7 +184,7 @@ def _build_tx(  # pylint: disable=unused-argument
                 amount=amount,
                 tx_fee=50000,
                 tx_nonce="0x",
-                chain_id=chain_type.id,
+                chain_id=chain_type.value,
                 raise_on_try=True,
             )
             return ledger_api.update_with_gas_estimate(