add support for OWNER_CHECK_721

Elnaril · May 24, 2024 · 178948b · 178948b
1 parent 51a8270
commit 178948b
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git a/uniswap_universal_router_decoder/_abi_builder.py b/uniswap_universal_router_decoder/_abi_builder.py
@@ -110,6 +110,7 @@ def build_abi_map(self) -> _ABIMap:
             _RouterFunction.TRANSFER: self._add_mapping(self._build_transfer),
             _RouterFunction.SEAPORT_V1_5: self._add_mapping(self._build_seaport_v1_5),
             _RouterFunction.SWEEP_ERC721: self._add_mapping(self._build_sweep_erc721),
+            _RouterFunction.OWNER_CHECK_721: self._add_mapping(self._build_owner_check_721),
         }
         return abi_map
 
@@ -186,3 +187,8 @@ def _build_seaport_v1_5() -> _FunctionABI:
     def _build_sweep_erc721() -> _FunctionABI:
         builder = _FunctionABIBuilder("SWEEP_ERC721")
         return builder.add_address("token").add_address("recipient").add_uint256("id").build()
+
+    @staticmethod
+    def _build_owner_check_721() -> _FunctionABI:
+        builder = _FunctionABIBuilder("OWNER_CHECK_721")
+        return builder.add_address("owner").add_address("token").add_uint256("id").build()
diff --git a/uniswap_universal_router_decoder/_encoder.py b/uniswap_universal_router_decoder/_encoder.py
@@ -635,6 +635,33 @@ def sweep_erc721(
         )
         return self
 
+    def _encode_owner_check_721_sub_contract(
+            self,
+            owner: ChecksumAddress,
+            token: ChecksumAddress,
+            token_id: int) -> HexStr:
+        abi_mapping = self._abi_map[_RouterFunction.OWNER_CHECK_721]
+        sub_contract = self._w3.eth.contract(abi=abi_mapping.fct_abi.get_full_abi())
+        contract_function: ContractFunction = sub_contract.functions.OWNER_CHECK_721(owner, token, token_id)
+        return remove_0x_prefix(encode_abi(self._w3, contract_function.abi, [owner, token, token_id]))
+
+    def owner_check_721(
+            self,
+            owner_address: ChecksumAddress,
+            token_address: ChecksumAddress,
+            token_id: int) -> _ChainedFunctionBuilder:
+        self.commands.append(_RouterFunction.OWNER_CHECK_721.value)
+        self.arguments.append(
+            Web3.to_bytes(
+                hexstr=self._encode_owner_check_721_sub_contract(
+                    owner_address,
+                    token_address,
+                    token_id,
+                )
+            )
+        )
+        return self
+
     def build(self, deadline: Optional[int] = None) -> HexStr:
         """
         Build the encoded input for all the chained commands, ready to be sent to the UR

diff --git a/uniswap_universal_router_decoder/_enums.py b/uniswap_universal_router_decoder/_enums.py
@@ -26,6 +26,7 @@ class _RouterFunction(Enum):
     WRAP_ETH = 11
     UNWRAP_WETH = 12
     SEAPORT_V1_5 = 16
+    OWNER_CHECK_721 = 21
     SWEEP_ERC721 = 23