From a96f0c5be765ab925ce9fe572199eb21c6bfc62c Mon Sep 17 00:00:00 2001
From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com>
Date: Sat, 24 Feb 2024 04:12:20 -0500
Subject: [PATCH] feat: Contract class (#124)

* feat: patched brownie.Contract class

* fix: patch_contract
---
 dank_mids/__init__.py               |  2 +-
 dank_mids/brownie_patch/__init__.py |  2 +-
 dank_mids/brownie_patch/contract.py | 35 ++++++++++++++++++++---------
 3 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/dank_mids/__init__.py b/dank_mids/__init__.py
index 5b253259..236cf42a 100644
--- a/dank_mids/__init__.py
+++ b/dank_mids/__init__.py
@@ -4,7 +4,7 @@
 from dank_mids.middleware import dank_middleware
 
 try:
-    from dank_mids.brownie_patch import dank_web3, patch_contract
+    from dank_mids.brownie_patch import Contract, dank_web3, patch_contract
 except ImportError:
     pass
 
diff --git a/dank_mids/brownie_patch/__init__.py b/dank_mids/brownie_patch/__init__.py
index f72846ed..81f24939 100644
--- a/dank_mids/brownie_patch/__init__.py
+++ b/dank_mids/brownie_patch/__init__.py
@@ -5,7 +5,7 @@
 try:
     from brownie import network, web3
     if network.is_connected():
-        from dank_mids.brownie_patch.contract import patch_contract
+        from dank_mids.brownie_patch.contract import Contract, patch_contract
         dank_web3 = setup_dank_w3_from_sync(web3)
 except ImportError:
     pass
\ No newline at end of file
diff --git a/dank_mids/brownie_patch/contract.py b/dank_mids/brownie_patch/contract.py
index b3a2838c..ae60bd37 100644
--- a/dank_mids/brownie_patch/contract.py
+++ b/dank_mids/brownie_patch/contract.py
@@ -1,14 +1,37 @@
 
 from types import MethodType
-from typing import Optional, Union
+from typing import Optional, Union, overload
 
-from brownie import Contract, network
+import brownie
 from brownie.network.contract import ContractCall, ContractTx, OverloadedMethod
 from web3 import Web3
 
 from dank_mids.brownie_patch.call import _get_coroutine_fn
 from dank_mids.brownie_patch.overloaded import _patch_overloaded_method
 
+
+class Contract(brownie.Contract):
+    """a modified `brownie.Contract` with async and call batching functionalities"""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        patch_contract(self)
+
+@overload
+def patch_contract(contract: Contract, w3: Optional[Web3] = None) -> Contract:...
+@overload
+def patch_contract(contract: Union[brownie.Contract, str], w3: Optional[Web3] = None) -> brownie.Contract:...
+def patch_contract(contract: Union[Contract, brownie.Contract, str], w3: Optional[Web3] = None) -> Union[Contract, brownie.Contract]:
+    """returns a patched version of `contract` with async and call batchings functionalities"""
+    if not isinstance(contract, brownie.Contract):
+        contract = brownie.Contract(contract)
+    if w3 is None and brownie.network.is_connected():
+        from dank_mids import dank_web3 as w3
+    if w3 is None:
+        raise RuntimeError("You must make sure either brownie is connected or you pass in a Web3 instance.")
+    for k, v in contract.__dict__.items():
+        _patch_if_method(v, w3)
+    return contract
+
 ContractMethod = Union[ContractCall, ContractTx, OverloadedMethod]
 
 def _patch_if_method(method: ContractMethod, w3: Web3) -> None:
@@ -16,11 +39,3 @@ def _patch_if_method(method: ContractMethod, w3: Web3) -> None:
         method.coroutine = MethodType(_get_coroutine_fn(w3, len(method.abi['inputs'])), method)
     elif isinstance(method, OverloadedMethod):
         _patch_overloaded_method(method, w3)
-
-def patch_contract(contract: Contract, w3: Optional[Web3] = None) -> Contract:
-    if w3 is None:
-        assert network.is_connected(), "You must make sure either brownie is connected or you pass in a Web3 instance."
-        from dank_mids.helpers import dank_web3 as w3
-    for k, v in contract.__dict__.items():
-        _patch_if_method(v, w3)
-    return contract