prediction-market-agent-tooling 0.65.5__py3-none-any.whl → 0.69.17.dev1149__py3-none-any.whl

prediction_market_agent_tooling/tools/contract.py

@@ -5,19 +5,27 @@ import typing as t
 from contextlib import contextmanager
 
 import eth_abi
+import tenacity
 from eth_abi.exceptions import DecodingError
 from pydantic import BaseModel, field_validator
 from web3 import Web3
-from web3.constants import CHECKSUM_ADDRESSS_ZERO
+from web3.constants import CHECKSUM_ADDRESSS_ZERO, HASH_ZERO
 from web3.contract.contract import Contract as Web3Contract
+from web3.exceptions import ContractCustomError, ContractLogicError
+from web3.types import BlockIdentifier
 
+from prediction_market_agent_tooling.chains import POLYGON_CHAIN_ID
 from prediction_market_agent_tooling.config import APIKeys, RPCConfig
 from prediction_market_agent_tooling.gtypes import (
     ABI,
     ChainID,
     ChecksumAddress,
     CollateralToken,
+    HexAddress,
+    HexBytes,
+    HexStr,
     Nonce,
+    OutcomeWei,
     TxParams,
     TxReceipt,
     Wei,
@@ -90,6 +98,7 @@ class ContractBaseClass(BaseModel):
         function_name: str,
         function_params: t.Optional[list[t.Any] | dict[str, t.Any]] = None,
         web3: Web3 | None = None,
+        block_identifier: BlockIdentifier | None = None,
     ) -> t.Any:
         """
         Used for reading from the contract.
@@ -101,6 +110,7 @@ class ContractBaseClass(BaseModel):
             contract_abi=self.abi,
             function_name=function_name,
             function_params=function_params,
+            block_identifier=block_identifier,
         )
 
     def send(
@@ -111,6 +121,7 @@ class ContractBaseClass(BaseModel):
         tx_params: t.Optional[TxParams] = None,
         timeout: int = 180,
         web3: Web3 | None = None,
+        default_gas: int | None = None,
     ) -> TxReceipt:
         """
         Used for changing a state (writing) to the contract.
@@ -127,6 +138,7 @@ class ContractBaseClass(BaseModel):
                 function_params=function_params,
                 tx_params=tx_params,
                 timeout=timeout,
+                default_gas=default_gas,
             )
         return send_function_on_contract_tx(
             web3=web3 or self.get_web3(),
@@ -137,6 +149,7 @@ class ContractBaseClass(BaseModel):
             function_params=function_params,
             tx_params=tx_params,
             timeout=timeout,
+            default_gas=default_gas,
         )
 
     def send_with_value(
@@ -186,6 +199,22 @@ class ContractProxyBaseClass(ContractBaseClass):
         return Web3.to_checksum_address(address)
 
 
+class ContractProcessorBaseClass(ContractBaseClass):
+    """
+    Contract base class for processor contracts.
+    """
+
+    abi: ABI = abi_field_validator(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "../abis/processor.abi.json"
+        )
+    )
+
+    def processor(self, web3: Web3 | None = None) -> ChecksumAddress:
+        address = self.call("processor", web3=web3)
+        return Web3.to_checksum_address(address)
+
+
 class ContractERC20BaseClass(ContractBaseClass):
     """
     Contract base class extended by ERC-20 standard methods.
@@ -197,10 +226,18 @@ class ContractERC20BaseClass(ContractBaseClass):
         )
     )
 
+    def name(self, web3: Web3 | None = None) -> str:
+        name: str = self.call("name", web3=web3)
+        return name
+
     def symbol(self, web3: Web3 | None = None) -> str:
         symbol: str = self.call("symbol", web3=web3)
         return symbol
 
+    def decimals(self, web3: Web3 | None = None) -> int:
+        decimals = int(self.call("decimals", web3=web3))
+        return decimals
+
     def symbol_cached(self, web3: Web3 | None = None) -> str:
         web3 = web3 or self.get_web3()
         cache_key = create_contract_method_cache_key(self.symbol, web3)
@@ -256,8 +293,17 @@ class ContractERC20BaseClass(ContractBaseClass):
             web3=web3,
         )
 
-    def balanceOf(self, for_address: ChecksumAddress, web3: Web3 | None = None) -> Wei:
-        balance = Wei(self.call("balanceOf", [for_address], web3=web3))
+    def balanceOf(
+        self,
+        for_address: ChecksumAddress,
+        web3: Web3 | None = None,
+        block_identifier: BlockIdentifier | None = None,
+    ) -> Wei:
+        balance = Wei(
+            self.call(
+                "balanceOf", [for_address], web3=web3, block_identifier=block_identifier
+            )
+        )
         return balance
 
     def balance_of_in_tokens(
@@ -413,6 +459,39 @@ class ContractERC4626BaseClass(ContractERC20BaseClass):
         return self.convertToShares(amount, web3=web3)
 
 
+class ContractWrapped1155BaseClass(ContractERC20BaseClass):
+    """Wrapped 1155 contract from Seer (https://gnosisscan.io/address/0x2f9c49974ad8b9b31424d9dc812667b16310ca50#readContract)
+    Source code - https://github.com/seer-pm/demo/blob/main/contracts/src/interaction/1155-to-20/Wrapped1155Factory.sol#L224
+    This contract inherits from ERC20 and contains additional properties, such as multiToken (conditional Token contract implementation)
+    and tokenId (token identifier). Goal is to wrap individual tokens into a standalone ERC20 token.
+    """
+
+    abi: ABI = abi_field_validator(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "../abis/erc1155.abi.json"
+        )
+    )
+
+    def factory(self, web3: Web3 | None = None) -> ChecksumAddress:
+        return Web3.to_checksum_address(self.call("factory", web3=web3))
+
+    def mint(
+        self,
+        api_keys: APIKeys,
+        to_address: ChecksumAddress,
+        amount: Wei,
+        tx_params: t.Optional[TxParams] = None,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        return self.send(
+            api_keys=api_keys,
+            function_name="mint",
+            function_params=[to_address, amount],
+            tx_params=tx_params,
+            web3=web3,
+        )
+
+
 class OwnableContract(ContractBaseClass):
     abi: ABI = abi_field_validator(
         os.path.join(
@@ -502,9 +581,19 @@ class ContractOnGnosisChain(ContractBaseClass):
     Contract base class with Gnosis Chain configuration.
     """
 
+    # This is defined like so because other chains (like Ethereum) rely on contracts that inherit
+    # from ContractOnGnosisChain. To be re-evaluated on https://github.com/gnosis/prediction-market-agent-tooling/issues/845
     CHAIN_ID = RPCConfig().chain_id
 
 
+class ContractOnPolygonChain(ContractBaseClass):
+    """
+    Contract base class with Gnosis Chain configuration.
+    """
+
+    CHAIN_ID = POLYGON_CHAIN_ID
+
+
 class ContractProxyOnGnosisChain(ContractProxyBaseClass, ContractOnGnosisChain):
     """
     Proxy contract base class with Gnosis Chain configuration.
@@ -533,6 +622,12 @@ class ContractDepositableWrapperERC20OnGnosisChain(
     """
 
 
+class ContractWrapped1155OnGnosisChain(
+    ContractWrapped1155BaseClass, ContractERC20OnGnosisChain
+):
+    pass
+
+
 class ContractERC4626OnGnosisChain(
     ContractERC4626BaseClass, ContractERC20OnGnosisChain
 ):
@@ -546,6 +641,254 @@ class ContractERC4626OnGnosisChain(
         return to_gnosis_chain_contract(super().get_asset_token_contract(web3=web3))
 
 
+class PayoutRedemptionEvent(BaseModel):
+    redeemer: HexAddress
+    collateralToken: HexAddress
+    parentCollectionId: HexBytes
+    conditionId: HexBytes
+    indexSets: list[int]
+    payout: Wei
+
+
+class ConditionPreparationEvent(BaseModel):
+    conditionId: HexBytes
+    oracle: HexAddress
+    questionId: HexBytes
+    outcomeSlotCount: int
+
+
+class ConditionalTokenContract(ContractBaseClass):
+    # Contract ABI taken from https://gnosisscan.io/address/0xCeAfDD6bc0bEF976fdCd1112955828E00543c0Ce#code.
+    abi: ABI = abi_field_validator(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)),
+            "../abis/omen_fpmm_conditionaltokens.abi.json",
+        )
+    )
+
+    def getConditionId(
+        self,
+        question_id: HexBytes,
+        oracle_address: ChecksumAddress,
+        outcomes_slot_count: int,
+        web3: Web3 | None = None,
+    ) -> HexBytes:
+        id_ = HexBytes(
+            self.call(
+                "getConditionId",
+                [oracle_address, question_id, outcomes_slot_count],
+                web3=web3,
+            )
+        )
+        return id_
+
+    def balanceOf(
+        self, from_address: ChecksumAddress, position_id: int, web3: Web3 | None = None
+    ) -> OutcomeWei:
+        balance = OutcomeWei(
+            self.call("balanceOf", [from_address, position_id], web3=web3)
+        )
+        return balance
+
+    def isApprovedForAll(
+        self,
+        owner: ChecksumAddress,
+        for_address: ChecksumAddress,
+        web3: Web3 | None = None,
+    ) -> bool:
+        is_approved: bool = self.call(
+            "isApprovedForAll", [owner, for_address], web3=web3
+        )
+        return is_approved
+
+    def getCollectionId(
+        self,
+        parent_collection_id: HexBytes,
+        condition_id: HexBytes,
+        index_set: int,
+        web3: Web3 | None = None,
+    ) -> HexBytes:
+        collection_id = HexBytes(
+            self.call(
+                "getCollectionId",
+                [parent_collection_id, condition_id, index_set],
+                web3=web3,
+            )
+        )
+        return collection_id
+
+    def getPositionId(
+        self,
+        collateral_token_address: ChecksumAddress,
+        collection_id: HexBytes,
+        web3: Web3 | None = None,
+    ) -> int:
+        position_id: int = self.call(
+            "getPositionId",
+            [collateral_token_address, collection_id],
+            web3=web3,
+        )
+        return position_id
+
+    def mergePositions(
+        self,
+        api_keys: APIKeys,
+        collateral_token_address: ChecksumAddress,
+        conditionId: HexBytes,
+        index_sets: t.List[int],
+        amount: OutcomeWei,
+        parent_collection_id: HexStr = HASH_ZERO,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        return self.send(
+            api_keys=api_keys,
+            function_name="mergePositions",
+            function_params=[
+                collateral_token_address,
+                parent_collection_id,
+                conditionId,
+                index_sets,
+                amount,
+            ],
+            web3=web3,
+        )
+
+    def redeemPositions(
+        self,
+        api_keys: APIKeys,
+        collateral_token_address: HexAddress,
+        condition_id: HexBytes,
+        index_sets: t.List[int],
+        parent_collection_id: HexStr = HASH_ZERO,
+        web3: Web3 | None = None,
+    ) -> PayoutRedemptionEvent:
+        receipt_tx = self.send(
+            api_keys=api_keys,
+            function_name="redeemPositions",
+            function_params=[
+                collateral_token_address,
+                parent_collection_id,
+                condition_id,
+                index_sets,
+            ],
+            web3=web3,
+        )
+        redeem_event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.PayoutRedemption()
+            .process_receipt(receipt_tx)
+        )
+        logger.info(
+            f"Receipt tx: `{receipt_tx}` Redeem event logs: `{redeem_event_logs}`"
+        )
+        redeem_event = PayoutRedemptionEvent(**redeem_event_logs[0]["args"])
+        return redeem_event
+
+    def getOutcomeSlotCount(
+        self, condition_id: HexBytes, web3: Web3 | None = None
+    ) -> int:
+        count: int = self.call("getOutcomeSlotCount", [condition_id], web3=web3)
+        return count
+
+    def does_condition_exists(
+        self, condition_id: HexBytes, web3: Web3 | None = None
+    ) -> bool:
+        return self.getOutcomeSlotCount(condition_id, web3=web3) > 0
+
+    def is_condition_resolved(
+        self, condition_id: HexBytes, web3: Web3 | None = None
+    ) -> bool:
+        # from ConditionalTokens.redeemPositions:
+        # uint den = payoutDenominator[conditionId]; require(den > 0, "result for condition not received yet");
+        payout_for_condition = self.payoutDenominator(condition_id, web3=web3)
+        return payout_for_condition > 0
+
+    def payoutDenominator(
+        self, condition_id: HexBytes, web3: Web3 | None = None
+    ) -> int:
+        payoutForCondition: int = self.call(
+            "payoutDenominator", [condition_id], web3=web3
+        )
+        return payoutForCondition
+
+    def splitPosition(
+        self,
+        api_keys: APIKeys,
+        collateral_token: ChecksumAddress,
+        condition_id: HexBytes,
+        outcome_slot_count: int,
+        amount_wei: Wei,
+        parent_collateral_id: HexBytes = HexBytes(HASH_ZERO),
+        tx_params: t.Optional[TxParams] = None,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        # We always split the full set of outcome tokens (for simplicity)
+        # partitions are given in bitmasks, i.e. outcomes are 1,2,4,8,etc.
+        partition = [2**i for i in range(outcome_slot_count)]
+        return self.send(
+            api_keys=api_keys,
+            function_name="splitPosition",
+            function_params=[
+                collateral_token,
+                parent_collateral_id,
+                condition_id,
+                partition,
+                amount_wei,
+            ],
+            tx_params=tx_params,
+            web3=web3,
+        )
+
+    def setApprovalForAll(
+        self,
+        api_keys: APIKeys,
+        for_address: ChecksumAddress,
+        approve: bool,
+        tx_params: t.Optional[TxParams] = None,
+        web3: Web3 | None = None,
+    ) -> TxReceipt:
+        return self.send(
+            api_keys=api_keys,
+            function_name="setApprovalForAll",
+            function_params=[
+                for_address,
+                approve,
+            ],
+            tx_params=tx_params,
+            web3=web3,
+        )
+
+    def prepareCondition(
+        self,
+        api_keys: APIKeys,
+        oracle_address: ChecksumAddress,
+        question_id: HexBytes,
+        outcomes_slot_count: int,
+        tx_params: t.Optional[TxParams] = None,
+        web3: Web3 | None = None,
+    ) -> ConditionPreparationEvent:
+        receipt_tx = self.send(
+            api_keys=api_keys,
+            function_name="prepareCondition",
+            function_params=[
+                oracle_address,
+                question_id,
+                outcomes_slot_count,
+            ],
+            tx_params=tx_params,
+            web3=web3,
+        )
+
+        event_logs = (
+            self.get_web3_contract(web3=web3)
+            .events.ConditionPreparation()
+            .process_receipt(receipt_tx)
+        )
+        cond_event = ConditionPreparationEvent(**event_logs[0]["args"])
+
+        return cond_event
+
+
 class DebuggingContract(ContractOnGnosisChain):
     # Contract ABI taken from https://gnosisscan.io/address/0x5Aa82E068aE6a6a1C26c42E5a59520a74Cdb8998#code.
     abi: ABI = abi_field_validator(
@@ -594,64 +937,154 @@ def contract_implements_function(
     look_for_proxy_contract: bool = True,
 ) -> bool:
     function_signature = f"{function_name}({','.join(function_arg_types or [])})"
-    function_selector = web3.keccak(text=function_signature)[0:4].hex()[2:]
+    function_selector = web3.keccak(text=function_signature)[0:4].to_0x_hex()[2:]
     # 1. Check directly in bytecode
-    bytecode = web3.eth.get_code(contract_address).hex()
+    bytecode = web3.eth.get_code(contract_address).to_0x_hex()
     if function_selector in bytecode:
         return True
-    contract_code = web3.eth.get_code(contract_address).hex()
+    contract_code = web3.eth.get_code(contract_address).to_0x_hex()
     implements = function_selector in contract_code
 
     # If not found directly and we should check proxies
     if not implements and look_for_proxy_contract:
-        # Case 1: Check if it's a standard proxy (has implementation() function)
-        if contract_implements_function(
-            contract_address, "implementation", web3, look_for_proxy_contract=False
-        ):
-            # Get the implementation address and check the function there
-            implementation_address = ContractProxyOnGnosisChain(
-                address=contract_address
-            ).implementation()
+        imp_addresses = uni_implementation_address(contract_address, web3)
+        for imp_address in imp_addresses:
             implements = contract_implements_function(
-                implementation_address,
+                imp_address,
                 function_name=function_name,
                 web3=web3,
                 function_arg_types=function_arg_types,
                 look_for_proxy_contract=False,
             )
-        else:
-            # Case 2: Check if it's a minimal proxy contract
-            implements = minimal_proxy_implements_function(
-                contract_address=contract_address,
-                function_name=function_name,
-                web3=web3,
-                function_arg_types=function_arg_types,
-            )
+            # If one of the implementations has the function, we can terminate early.
+            if implements:
+                break
 
     return implements
 
 
-def minimal_proxy_implements_function(
-    contract_address: ChecksumAddress,
-    function_name: str,
-    web3: Web3,
-    function_arg_types: list[str] | None = None,
-) -> bool:
+def uni_implementation_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> list[ChecksumAddress]:
+    """
+    There are multiple ways how proxies can be implemented.
+    This function enumerates them and returns the ones that succeed, or an empty list.
+    """
+    # Currently, our implementation sometimes returns implementation multiple addresses for a single proxy contract.
+    # This should not happen and needs more investigation.
+    # Probably `minimal_proxy_address` is bugged and a function identifying the ERC-1967 proxy pattern (https://eips.ethereum.org/EIPS/eip-1967#logic-contract-address) should be implemented.
+    # TODO: Fix the above, and afterwards assert that only 1 imp address is returned from this function. Or prove that this could indeed happen (although we are very pretty sure it shouldn't).
+    addresses = [
+        implementation_proxy_address(contract_address, web3),
+        processor_proxy_address(contract_address, web3),
+        minimal_proxy_address(contract_address, web3),
+        seer_minimal_proxy_address(contract_address, web3),
+        eip_1967_proxy_address(contract_address, web3),
+        zeppelinos_unstructured_storage_proxy_address(contract_address, web3),
+    ]
+    return [addr for addr in addresses if addr is not None]
+
+
+def implementation_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
+    if not contract_implements_function(
+        contract_address, "implementation", web3, look_for_proxy_contract=False
+    ):
+        return None
+    try:
+        return ContractProxyOnGnosisChain(address=contract_address).implementation(web3)
+    except (ContractCustomError, ContractLogicError, tenacity.RetryError) as e:
+        if isinstance(e, tenacity.RetryError) and not isinstance(
+            e.last_attempt.exception(), (ContractCustomError, ContractLogicError)
+        ):
+            raise
+
+        # For example https://gnosisscan.io/address/0x3221a28ed2b2e955da64d1d299956f277562c95c#code,
+        # it has the implementation method, but it's only for admins.
+        logger.warning(
+            f"Failed to get implementation for {contract_address=} even though it has the method: {e}"
+        )
+        return None
+
+
+def processor_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
+    if not contract_implements_function(
+        contract_address, "processor", web3, look_for_proxy_contract=False
+    ):
+        return None
+    try:
+        return ContractProcessorBaseClass(address=contract_address).processor(web3)
+    except (ContractCustomError, ContractLogicError, tenacity.RetryError) as e:
+        if isinstance(e, tenacity.RetryError) and not isinstance(
+            e.last_attempt.exception(), (ContractCustomError, ContractLogicError)
+        ):
+            raise
+
+        # For example https://gnosisscan.io/address/0x3221a28ed2b2e955da64d1d299956f277562c95c#code,
+        # it has the processor method, but it's only for admins.
+        logger.warning(
+            f"Failed to get processor for {contract_address=} even though it has the method: {e}"
+        )
+        return None
+
+
+def minimal_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
     try:
         # Read storage slot 0 which should contain the implementation address in minimal proxies
         raw_slot_0 = web3.eth.get_storage_at(contract_address, 0)
         singleton_address = eth_abi.decode(["address"], raw_slot_0)[0]
-        # Recurse into singleton
-        return contract_implements_function(
-            Web3.to_checksum_address(singleton_address),
-            function_name=function_name,
-            web3=web3,
-            function_arg_types=function_arg_types,
-            look_for_proxy_contract=False,
-        )
+        return Web3.to_checksum_address(singleton_address)
     except DecodingError:
         logger.info(f"Error decoding contract address for singleton")
-        return False
+        return None
+
+
+def seer_minimal_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
+    try:
+        # Read address between specific indices to find logic contract
+        bytecode = web3.eth.get_code(contract_address)
+        logic_contract_address = bytecode[11:31]
+        if not Web3.is_address(logic_contract_address):
+            return None
+        return Web3.to_checksum_address(logic_contract_address)
+    except DecodingError:
+        logger.info("Error decoding contract address on seer minimal proxy")
+        return None
+
+
+def eip_1967_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
+    try:
+        slot = HexBytes(Web3.keccak(text="eip1967.proxy.implementation")).as_int() - 1
+        raw_slot = web3.eth.get_storage_at(contract_address, slot)
+        address = eth_abi.decode(["address"], raw_slot)[0]
+        return Web3.to_checksum_address(address)
+    except DecodingError:
+        logger.info("Error decoding contract address for eip 1967 proxy")
+        return None
+
+
+def zeppelinos_unstructured_storage_proxy_address(
+    contract_address: ChecksumAddress, web3: Web3
+) -> ChecksumAddress | None:
+    try:
+        slot = HexBytes(
+            Web3.keccak(text="org.zeppelinos.proxy.implementation")
+        ).as_int()
+        raw_slot = web3.eth.get_storage_at(contract_address, slot)
+        address = eth_abi.decode(["address"], raw_slot)[0]
+        return Web3.to_checksum_address(address)
+    except DecodingError:
+        logger.info("Error decoding contract address for zeppelinos proxy")
+        return None
 
 
 def init_collateral_token_contract(
@@ -673,6 +1106,13 @@ def init_collateral_token_contract(
     ):
         return ContractDepositableWrapperERC20BaseClass(address=address)
 
+    elif contract_implements_function(
+        address,
+        "multiToken",
+        web3=web3,
+    ):
+        return ContractWrapped1155BaseClass(address=address)
+
     elif contract_implements_function(
         address,
         "balanceOf",
@@ -694,6 +1134,8 @@ def to_gnosis_chain_contract(
         return ContractERC4626OnGnosisChain(address=contract.address)
     elif isinstance(contract, ContractDepositableWrapperERC20BaseClass):
         return ContractDepositableWrapperERC20OnGnosisChain(address=contract.address)
+    elif isinstance(contract, ContractWrapped1155BaseClass):
+        return ContractWrapped1155OnGnosisChain(address=contract.address)
     elif isinstance(contract, ContractERC20BaseClass):
         return ContractERC20OnGnosisChain(address=contract.address)
     else: