From 320af2aff5b1b4a25c8b689e3c61be3635d92329 Mon Sep 17 00:00:00 2001
From: BobTheBuidler <70677534+BobTheBuidler@users.noreply.github.com>
Date: Mon, 22 Apr 2024 17:03:07 -0400
Subject: [PATCH] feat(mypy): DankEth type checking helper (#180)

---
 dank_mids/helpers/_helpers.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/dank_mids/helpers/_helpers.py b/dank_mids/helpers/_helpers.py
index fc9eb007..0b124421 100644
--- a/dank_mids/helpers/_helpers.py
+++ b/dank_mids/helpers/_helpers.py
@@ -28,9 +28,14 @@
 T = TypeVar("T")
 P = ParamSpec("P")
 
+class DankEth(AsyncEth):
+    @alru_cache(ttl=0)
+    async def get_block_number(self) -> int:
+        return await super().get_block_number()
+    
 class DankWeb3:
-    """This is just a helper for type checkers. Your object will just be a modified `Web3` object."""
-    eth: AsyncEth
+    """This is just a helper for type checkers. Your object will just be a modified ``web3.Web3`` object."""
+    eth: DankEth
 
 def setup_dank_w3(async_w3: Web3) -> DankWeb3:
     """ Injects Dank Middleware into an async Web3 instance. """
@@ -42,11 +47,7 @@ def setup_dank_w3(async_w3: Web3) -> DankWeb3:
         async_w3.middleware_onion.inject(dank_middleware, layer=0)
         async_w3.middleware_onion.add(geth_poa_middleware)
         dank_w3s.append(async_w3)
-    async_w3.eth._get_block_number = async_w3.eth.get_block_number
-    @alru_cache(ttl=0)
-    async def get_block_number() -> int:
-        return await async_w3.eth._get_block_number()
-    async_w3.eth.get_block_number = get_block_number
+    async_w3.eth = DankEth(async_w3)
     return async_w3
 
 def setup_dank_w3_from_sync(sync_w3: Web3) -> DankWeb3: