pass RouterCodec w3 to _ABIBuilder

make abi_map built at init and available for *_v4_exact_input_params()
Elnaril · Feb 7, 2025 · d1dfd07 · d1dfd07
1 parent d03ecb4
commit d1dfd07
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/coverage.json b/coverage.json
diff --git a/uniswap_universal_router_decoder/_abi_builder.py b/uniswap_universal_router_decoder/_abi_builder.py
@@ -13,6 +13,7 @@
     cast,
     Dict,
     List,
+    Optional,
     Sequence,
     Union,
 )
@@ -150,8 +151,9 @@ def add_v4_exact_input_params(self, arg_name: str = "params") -> FunctionABIBuil
 
 
 class _ABIBuilder:
-    def __init__(self) -> None:
-        self.w3 = Web3()
+    def __init__(self, w3: Optional[Web3] = None) -> None:
+        self.w3 = w3 if w3 else Web3()
+        self.abi_map = self.build_abi_map()
         if not registry.has_encoder("ExactInputParams"):
             registry.register("ExactInputParams", self.encode_v4_exact_input_params, self.decode_v4_exact_input_params)
 
@@ -196,6 +198,7 @@ def build_abi_map(self) -> ABIMap:
             MiscFunctions.EXECUTE_WITH_DEADLINE: self._build_execute_with_deadline(),
             MiscFunctions.UNLOCK_DATA: self._build_unlock_data(),
             MiscFunctions.V4_POOL_ID: self._build_v4_pool_id(),
+            MiscFunctions.STRICT_V4_SWAP_EXACT_IN: self._build_strict_v4_swap_exact_in(),
         }
         return abi_map
 
@@ -425,13 +428,13 @@ def _build_v4_swap_exact_in() -> FunctionABI:
         return builder.add_v4_exact_input_params().build()
 
     def decode_v4_exact_input_params(self, stream: BytesIO) -> Dict[str, Any]:
-        fct_abi = self._build_strict_v4_swap_exact_in()
+        fct_abi = self.abi_map[MiscFunctions.STRICT_V4_SWAP_EXACT_IN]
         raw_data = stream.read()
         sub_contract = self.w3.eth.contract(abi=fct_abi.get_full_abi())
         fct_name, decoded_params = sub_contract.decode_function_input(fct_abi.get_selector() + raw_data[32:])
         return cast(Dict[str, Any], decoded_params)
 
     def encode_v4_exact_input_params(self, args: Sequence[Any]) -> bytes:
-        fct_abi = self._build_strict_v4_swap_exact_in()
+        fct_abi = self.abi_map[MiscFunctions.STRICT_V4_SWAP_EXACT_IN]
         encoded_data = 0x20.to_bytes(32, "big") + encode(fct_abi.get_abi_types(), args)
         return encoded_data
diff --git a/uniswap_universal_router_decoder/_enums.py b/uniswap_universal_router_decoder/_enums.py
@@ -100,3 +100,5 @@ class MiscFunctions(Enum):
     EXECUTE_WITH_DEADLINE = auto()  # value = "execute" would be nice, but enum names and values must be unique
     UNLOCK_DATA = auto()
     V4_POOL_ID = auto()
+
+    STRICT_V4_SWAP_EXACT_IN = auto()
diff --git a/uniswap_universal_router_decoder/router_codec.py b/uniswap_universal_router_decoder/router_codec.py
@@ -44,13 +44,13 @@
 
 class RouterCodec:
     def __init__(self, w3: Optional[Web3] = None, rpc_endpoint: Optional[str] = None) -> None:
-        self._abi_map = _ABIBuilder().build_abi_map()
         if w3:
             self._w3 = w3
         elif rpc_endpoint:
             self._w3 = Web3(Web3.HTTPProvider(rpc_endpoint))
         else:
             self._w3 = Web3()
+        self._abi_map = _ABIBuilder(self._w3).abi_map
         self.decode = _Decoder(self._w3, self._abi_map)
         self.encode = _Encoder(self._w3, self._abi_map)