Add classes to handle W3_RPC

.github/workflows/python_package.yaml

@@ -25,6 +25,7 @@ jobs:
       - name: Test with pytest
         env:
           ETHNODEURL: ${{ secrets.ETHNODEURL }}
+          DRPC_KEY: ${{ secrets.DRPC_KEY }}
         run: |
           pip install -r bal_tools/requirements-dev.txt
           pytest

bal_tools/__init__.py

@@ -7,3 +7,4 @@
 from .subgraph import Subgraph
 from .pools_gauges import BalPoolsGauges
 from .ecosystem import Aura
+from .drpc import Web3RpcByChain, Web3Rpc

bal_tools/drpc.py

@@ -0,0 +1,56 @@
+from web3 import Web3
+from bal_addresses import AddrBook
+
+DRPC_NAME_OVERRIDES = {
+    "mainnet": "ethereum",
+    "zkevm": "polygon-zkevm",
+}
+
+class Web3RpcByChain:
+    def __init__(self, DRPC_KEY):
+        self.DRPC_KEY = DRPC_KEY
+        self._w3_by_chain = {}
+
+    def __getitem__(self, chain):
+        return self._get_or_create_w3(chain)
+
+    def __getattr__(self, chain):
+        return self._get_or_create_w3(chain)
+
+    def _get_or_create_w3(self, chain):
+        if chain not in self._w3_by_chain:
+            w3 = Web3Rpc(chain, self.DRPC_KEY)
+            self._w3_by_chain[chain] = w3
+        return self._w3_by_chain[chain]
+
+    def __setitem__(self, chain, value):
+        self._w3_by_chain[chain] = value
+
+    def __delitem__(self, chain):
+        del self._w3_by_chain[chain]
+
+    def __iter__(self):
+        return iter(self._w3_by_chain)
+
+    def keys(self):
+        return self._w3_by_chain.keys()
+
+    def values(self):
+        return self._w3_by_chain.values()
+
+    def items(self):
+        return self._w3_by_chain.items()
+
+class Web3Rpc:
+    def __init__(self, chain, DRPC_KEY):
+        drpc_chain = DRPC_NAME_OVERRIDES.get(chain, chain)
+        self.w3 = Web3(Web3.HTTPProvider(f"https://lb.drpc.org/ogrpc?network={drpc_chain}&dkey={DRPC_KEY}"))
+        if not self.w3.is_connected():
+            raise ConnectionError(f"Unable to connect to {drpc_chain} network with provided DRPC_KEY.")
+        try:
+            self.w3.eth.block_number
+        except Exception as e:
+            raise ConnectionError(f"Error fetching latest block number: {e}, chain: {chain}, url: {self.w3.provider.endpoint_uri}")
+
+    def __getattr__(self, name):
+        return getattr(self.w3, name)

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-VERSION = "0.1.1"
+VERSION = "0.1.2"
 DESCRIPTION = "Balancer Tools"
 LONG_DESCRIPTION = "Balancer Maxi helper and ecosystem tools"
 

tests/conftest.py

@@ -91,3 +91,8 @@ def pool_snapshot_blocks():
         "avalanche": 30000000,
         "zkevm": 1200000,
     }
+
+
+@pytest.fixture(scope="module")
+def DRPC_KEY():
+    return os.getenv("DRPC_KEY")

tests/test_drpc.py

@@ -0,0 +1,27 @@
+import pytest
+from bal_tools.drpc import Web3RpcByChain, Web3Rpc
+from tests.conftest import chains
+
+
+def test_drpc_by_chain_dict_syntax(DRPC_KEY):
+    if not DRPC_KEY:
+        pytest.skip("Skipping DRPC_KEY not set")
+
+    w3_by_chain = Web3RpcByChain(DRPC_KEY)
+    assert not w3_by_chain.keys()
+
+    for chain in chains:
+        assert w3_by_chain[chain].eth
+
+    assert sorted(w3_by_chain.keys()) == sorted(chains)
+
+def test_drpc_by_chain_dot_syntax(DRPC_KEY):
+    if not DRPC_KEY:
+        pytest.skip("Skipping DRPC_KEY not set")
+
+    w3_by_chain = Web3RpcByChain(DRPC_KEY)
+
+    assert not w3_by_chain.keys()
+    assert w3_by_chain.mainnet.eth
+    assert w3_by_chain.arbitrum.eth
+    assert list(w3_by_chain.keys()) == ["mainnet", "arbitrum"]