feat: patch overloaded methods (#3)

dank_mids/brownie_patch/__init__.py

@@ -1,3 +1,2 @@
 
 from dank_mids.brownie_patch.contract import patch_contract
-

dank_mids/brownie_patch/account.py

@@ -5,15 +5,14 @@
 from types import MethodType
 from typing import Any, Dict, Optional, Tuple
 
-from web3 import Web3
-
 from brownie._config import CONFIG
 from brownie.convert import Wei
 from brownie.exceptions import VirtualMachineError
 from brownie.network.account import Account
 from brownie.network.rpc import Rpc
 from brownie.network.state import Chain
 from brownie.network.transaction import TransactionReceipt
+from web3 import Web3
 
 rpc = Rpc()
 

dank_mids/brownie_patch/call.py

@@ -4,12 +4,11 @@
 from typing import Coroutine, Dict, Tuple, Union
 
 import eth_abi
-from hexbytes import HexBytes
-from web3 import Web3
-
 from brownie.exceptions import VirtualMachineError
 from brownie.network.contract import ContractCall, _get_tx
 from brownie.project.compiler.solidity import SOLIDITY_ERROR_CODES
+from hexbytes import HexBytes
+from web3 import Web3
 
 
 def _patch_call(call: ContractCall, w3: Web3) -> None:

dank_mids/brownie_patch/contract.py

@@ -1,22 +1,29 @@
 
-from typing import Optional
+from typing import Optional, Union
 
+from brownie import Contract, network, web3
+from brownie.network.contract import ContractCall, ContractTx, OverloadedMethod
 from dank_mids.brownie_patch.call import _patch_call
+from dank_mids.brownie_patch.overloaded import _patch_overloaded_method
 from dank_mids.brownie_patch.tx import _patch_tx
 from web3 import Web3
 
-from brownie import Contract, network, web3
-from brownie.network.contract import ContractCall, ContractTx
+ContractMethod = Union[ContractCall, ContractTx, OverloadedMethod]
+
+def _patch_if_method(method: ContractMethod, w3: Web3):
+    if isinstance(method, ContractCall) or isinstance(method, ContractTx):
+        _patch_call(method, w3)
+    # TODO implement this properly
+    #elif isinstance(method, ContractTx):
+    #    _patch_tx(method, w3)
+    elif isinstance(method, OverloadedMethod):
+        _patch_overloaded_method(method, w3)
 
 
 def patch_contract(contract: Contract, w3: Optional[Web3]) -> Contract:
     if w3 is None and network.is_connected():
         w3 = web3
     assert w3 is not None, "You must make sure either brownie is connected or you pass in a Web3 instance."
     for k, v in contract.__dict__.items():
-        if isinstance(v, ContractCall):
-            _patch_call(v, w3)
-        # TODO implement this properly
-        #elif isinstance(v, ContractTx):
-        #    _patch_tx(v, w3)
+        _patch_if_method(v, w3)
     return contract

dank_mids/brownie_patch/overloaded.py

@@ -0,0 +1,31 @@
+import functools
+from types import MethodType
+from typing import Coroutine, Dict, Tuple, Union
+
+from brownie.network.contract import ContractCall, ContractTx, OverloadedMethod
+from dank_mids.brownie_patch.call import _patch_call
+from dank_mids.brownie_patch.tx import _patch_tx
+from web3 import Web3
+
+
+def _patch_overloaded_method(call: OverloadedMethod, w3: Web3) -> None:
+    @functools.wraps(call)
+    async def coroutine(
+        self,
+        *args: Tuple,
+        block_identifier: Union[int, str, bytes] = None,
+        override: Dict = None
+    ) -> Coroutine:
+        fn = self._get_fn_from_args(args)
+        kwargs = {"block_identifier": block_identifier, "override": override}
+        kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        return await fn.coroutine(*args, **kwargs)
+
+    for args, method in call.__dict__['methods'].items():
+        if isinstance(method, ContractCall) or isinstance(method, ContractTx):
+            _patch_call(method, w3)
+    # TODO implement this properly
+        #elif isinstance(call, ContractTx):
+            #_patch_tx(call, w3)
+
+    call.coroutine = MethodType(coroutine, call)

dank_mids/brownie_patch/tx.py

@@ -4,9 +4,9 @@
 from typing import Coroutine, Dict, Tuple, Union
 
 from brownie.network.contract import ContractTx, _get_tx
+from dank_mids.brownie_patch.account import _patch_account
 from web3 import Web3
 
-from dank_mids.brownie_patch.account import _patch_account
 
 def _patch_tx(call: ContractTx, w3: Web3) -> None:
 

tests/test_brownie_patch.py

@@ -1,7 +1,7 @@
 
 from brownie import Contract
 from dank_mids.brownie_patch import patch_contract
-from dank_mids.brownie_patch.call import _patch_call
+from dank_mids.brownie_patch.contractcall import _patch_call
 from dank_mids.brownie_patch.tx import _patch_tx
 from multicall.utils import await_awaitable