From b712f23212a59f6a37d45ad30cbf1575ef581d0a Mon Sep 17 00:00:00 2001 From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com> Date: Wed, 18 May 2022 02:17:56 -0400 Subject: [PATCH] feat: patch overloaded methods (#3) --- dank_mids/brownie_patch/__init__.py | 1 - dank_mids/brownie_patch/account.py | 3 +-- dank_mids/brownie_patch/call.py | 5 ++--- dank_mids/brownie_patch/contract.py | 23 +++++++++++++------- dank_mids/brownie_patch/overloaded.py | 31 +++++++++++++++++++++++++++ dank_mids/brownie_patch/tx.py | 2 +- tests/test_brownie_patch.py | 2 +- 7 files changed, 51 insertions(+), 16 deletions(-) create mode 100644 dank_mids/brownie_patch/overloaded.py diff --git a/dank_mids/brownie_patch/__init__.py b/dank_mids/brownie_patch/__init__.py index 9796db75..0ef4767c 100644 --- a/dank_mids/brownie_patch/__init__.py +++ b/dank_mids/brownie_patch/__init__.py @@ -1,3 +1,2 @@ from dank_mids.brownie_patch.contract import patch_contract - diff --git a/dank_mids/brownie_patch/account.py b/dank_mids/brownie_patch/account.py index 0af8d66b..3f3c3f51 100644 --- a/dank_mids/brownie_patch/account.py +++ b/dank_mids/brownie_patch/account.py @@ -5,8 +5,6 @@ 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 @@ -14,6 +12,7 @@ from brownie.network.rpc import Rpc from brownie.network.state import Chain from brownie.network.transaction import TransactionReceipt +from web3 import Web3 rpc = Rpc() diff --git a/dank_mids/brownie_patch/call.py b/dank_mids/brownie_patch/call.py index e6a7dbc6..3ed897fb 100644 --- a/dank_mids/brownie_patch/call.py +++ b/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: diff --git a/dank_mids/brownie_patch/contract.py b/dank_mids/brownie_patch/contract.py index 82025029..2970a571 100644 --- a/dank_mids/brownie_patch/contract.py +++ b/dank_mids/brownie_patch/contract.py @@ -1,12 +1,23 @@ -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: @@ -14,9 +25,5 @@ def patch_contract(contract: Contract, w3: Optional[Web3]) -> Contract: 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 diff --git a/dank_mids/brownie_patch/overloaded.py b/dank_mids/brownie_patch/overloaded.py new file mode 100644 index 00000000..576bb313 --- /dev/null +++ b/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) diff --git a/dank_mids/brownie_patch/tx.py b/dank_mids/brownie_patch/tx.py index 58b3a247..4cb6b6ac 100644 --- a/dank_mids/brownie_patch/tx.py +++ b/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: diff --git a/tests/test_brownie_patch.py b/tests/test_brownie_patch.py index c985d713..d32d852a 100644 --- a/tests/test_brownie_patch.py +++ b/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