feat: include proxy ABIs

src/ape/managers/chain.py

@@ -798,26 +798,48 @@ def cache_deployment(self, contract_instance: ContractInstance):
         """
 
         address = contract_instance.address
-        contract_type = contract_instance.contract_type
+        contract_type = contract_instance.contract_type  # may be a proxy
 
         # Cache contract type in memory before proxy check,
         # in case it is needed somewhere. It may get overridden.
         self._local_contract_types[address] = contract_type
 
-        proxy_info = self.provider.network.ecosystem.get_proxy_info(address)
-        if proxy_info:
+        if proxy_info := self.provider.network.ecosystem.get_proxy_info(address):
+            # The user is caching a deployment of a proxy with the target already set.
             self.cache_proxy_info(address, proxy_info)
-            contract_type = self.get(proxy_info.target) or contract_type
-            if contract_type:
+            if implementation_contract := self.get(proxy_info.target):
+                # Include proxy ABIs but ignore fallback/ctor etc.
+                proxy_abis = [
+                    abi for abi in contract_type.abi if abi.type in ("error", "event", "function")
+                ]
+                implementation_contract.abi.extend(proxy_abis)
+
+                # Include "hidden" ABIs, such as Safe's `masterCopy()`.
+                if proxy_info.abi and proxy_info.abi.signature not in [
+                    abi.signature for abi in contract_type.abi
+                ]:
+                    implementation_contract.abi.append(proxy_info.abi)
+
+                self._cache_contract_type(address, implementation_contract)
+
+                # Use this contract type in the user's contract instance.
+                contract_instance.contract_type = implementation_contract
+
+            else:
+                # No implementation yet. Just cache proxy.
                 self._cache_contract_type(address, contract_type)
 
-            return
+        else:
+            # Regular contract. Cache normally.
+            self._cache_contract_type(address, contract_type)
 
+        # Cache the deployment now.
         txn_hash = contract_instance.txn_hash
-        self._cache_contract_type(address, contract_type)
         if contract_type.name:
             self._cache_deployment(address, contract_type, txn_hash)
 
+        return contract_type
+
     def cache_proxy_info(self, address: AddressType, proxy_info: ProxyInfoAPI):
         """
         Cache proxy info for a particular address, useful for plugins adding already
@@ -1261,6 +1283,12 @@ def instance_at(
                     txn_hash = deployment["transaction_hash"]
                     break
 
+        # Include proxy-related ABIs.
+        if info := self.get_proxy_info(contract_address):
+            if proxy_abi := info.abi:
+                if proxy_abi.signature not in [x.signature for x in contract_type.abi]:
+                    contract_type.abi.append(proxy_abi)
+
         return ContractInstance(contract_address, contract_type, txn_hash=txn_hash)
 
     def instance_from_receipt(
@@ -1271,6 +1299,7 @@ def instance_from_receipt(
 
         Args:
             receipt (:class:`~ape.api.transactions.ReceiptAPI`): The receipt.
+            contract_type (ContractType): The deployed contract type.
 
         Returns:
             :class:`~ape.contracts.base.ContractInstance`

src/ape_ethereum/ecosystem.py

@@ -462,8 +462,7 @@ def get_proxy_info(self, address: AddressType) -> Optional[ProxyInfo]:
         if isinstance(contract_code, bytes):
             contract_code = to_hex(contract_code)
 
-        code = contract_code[2:]
-        if not code:
+        if not (code := contract_code[2:]):
             return None
 
         patterns = {
@@ -515,7 +514,7 @@ def str_to_slot(text):
             if _type == ProxyType.Beacon:
                 target = ContractCall(IMPLEMENTATION_ABI, target)(skip_trace=True)
 
-            return ProxyInfo(type=_type, target=target)
+            return ProxyInfo(type=_type, target=target, abi=MASTER_COPY_ABI)
 
         # safe >=1.1.0 provides `masterCopy()`, which is also stored in slot 0
         # call it and check that target matches
@@ -525,7 +524,8 @@ def str_to_slot(text):
             target = self.conversion_manager.convert(slot_0[-20:], AddressType)
             # NOTE: `target` is set in initialized proxies
             if target != ZERO_ADDRESS and target == singleton:
-                return ProxyInfo(type=ProxyType.GnosisSafe, target=target)
+                return ProxyInfo(type=ProxyType.GnosisSafe, target=target, abi=MASTER_COPY_ABI)
+
         except ApeException:
             pass
 
@@ -541,7 +541,7 @@ def str_to_slot(text):
                 target = ContractCall(IMPLEMENTATION_ABI, address)(skip_trace=True)
                 # avoid recursion
                 if target != ZERO_ADDRESS:
-                    return ProxyInfo(type=ProxyType.Delegate, target=target)
+                    return ProxyInfo(type=ProxyType.Delegate, target=target, abi=IMPLEMENTATION_ABI)
 
             except (ApeException, ValueError):
                 pass

src/ape_ethereum/proxies.py

@@ -1,5 +1,5 @@
 from enum import IntEnum, auto
-from typing import cast
+from typing import Optional, cast
 
 from eth_pydantic_types.hex import HexStr
 from ethpm_types import ContractType, MethodABI
@@ -69,6 +69,15 @@ class ProxyType(IntEnum):
 class ProxyInfo(ProxyInfoAPI):
     type: ProxyType
 
+    def __init__(self, **kwargs):
+        abi = kwargs.pop("abi", None)
+        super().__init__(**kwargs)
+        self._abi = abi
+
+    @property
+    def abi(self) -> Optional[MethodABI]:
+        return self._abi
+
 
 MASTER_COPY_ABI = MethodABI(
     type="function",

tests/functional/geth/test_proxy.py

@@ -57,20 +57,40 @@ def test_uups_proxy(get_contract_type, geth_contract, owner, ethereum):
 
 
 @geth_process_test
-def test_gnosis_safe(get_contract_type, geth_contract, owner, ethereum):
+def test_gnosis_safe(get_contract_type, geth_contract, owner, ethereum, chain):
+    # Setup a proxy contract.
     _type = get_contract_type("SafeProxy")
     contract = ContractContainer(_type)
-
     target = geth_contract.address
+    proxy_instance = owner.deploy(contract, target)
 
-    contract_instance = owner.deploy(contract, target)
-
-    actual = ethereum.get_proxy_info(contract_instance.address)
-
+    # (test)
+    actual = ethereum.get_proxy_info(proxy_instance.address)
     assert actual is not None
     assert actual.type == ProxyType.GnosisSafe
     assert actual.target == target
 
+    # Ensure we can call the proxy-method.
+    assert proxy_instance.masterCopy()
+
+    # Ensure we can target methods.
+    assert isinstance(proxy_instance.myNumber(), int)
+
+    # Ensure this works with new instances.
+    proxy_instance_ref_2 = chain.contracts.instance_at(proxy_instance.address)
+    assert proxy_instance_ref_2.masterCopy()
+    assert isinstance(proxy_instance_ref_2.myNumber(), int)
+
+    # Same - but clear the proxy ABI from the cached type.
+    chain.contracts._local_contract_types[proxy_instance.address].abi = [
+        x
+        for x in chain.contracts._local_contract_types[proxy_instance.address].abi
+        if x.type == "function" and x.name != "masterCopy"
+    ]
+    proxy_instance_ref_3 = chain.contracts.instance_at(proxy_instance.address)
+    assert proxy_instance_ref_3.masterCopy()
+    assert isinstance(proxy_instance_ref_3.myNumber(), int)
+
 
 @geth_process_test
 def test_openzeppelin(get_contract_type, geth_contract, owner, ethereum, sender):