fix: import error on web3.py < 6.7.0

banteg · Nov 27, 2024 · a64ed2f · a64ed2f
1 parent 3eb7641
commit a64ed2f
Show file tree

Hide file tree

Showing 3 changed files with 1,884 additions and 1,484 deletions.
diff --git a/multicall/utils.py b/multicall/utils.py
@@ -6,17 +6,22 @@
 
 import eth_retry
 from aiohttp import ClientTimeout
-from web3 import AsyncHTTPProvider, Web3, WebsocketProviderV2
+from web3 import AsyncHTTPProvider, Web3
 from web3.eth import AsyncEth
 from web3.providers.async_base import AsyncBaseProvider
 
 from multicall.constants import (AIOHTTP_TIMEOUT, ASYNC_SEMAPHORE, 
                                  NO_STATE_OVERRIDE, NUM_PROCESSES)
 
 try:
-    from web3 import AsyncWeb3
+    from web3 import AsyncWeb3, WebsocketProviderV2
 except ImportError:
-    AsyncWeb3 = None
+    AsyncWeb3, WebsocketProviderV2 = None, None
+
+try:
+    from web3 import WebsocketProviderV2
+except ImportError:
+    WebsocketProviderV2 = None
 
 
 chainids: Dict[Web3,int] = {}
@@ -59,15 +64,15 @@ def get_async_w3(w3: Web3) -> Web3:
 
     endpoint = get_endpoint(w3)
     request_kwargs = {'timeout': AIOHTTP_TIMEOUT}
-    if endpoint.startswith(("wss:", "ws:")):
+    if WebsocketProviderV2 and endpoint.startswith(("wss:", "ws:")):
         provider = WebsocketProviderV2(endpoint, request_kwargs)
     else:
         provider=AsyncHTTPProvider(endpoint, request_kwargs)
 
     # In older web3 versions, AsyncHTTPProvider objects come
     # with incompatible synchronous middlewares by default.
     middlewares = []
-    if AsyncWeb3:
+    if AsyncWeb3 and WebsocketProviderV2:
         async_w3 = AsyncWeb3(provider=provider, middlewares=middlewares)
     else:
         async_w3 = Web3(provider=provider, middlewares=middlewares)