From eda83eb3533e149c89d6b88520eb53f6b1a7a5b1 Mon Sep 17 00:00:00 2001
From: Juliya Smith <jules@apeworx.io>
Date: Mon, 2 Dec 2024 16:16:12 -0600
Subject: [PATCH] feat: make work on both

---
 src/ape/utils/_web3_compat.py          | 22 ++++++++++++++++++++++
 src/ape_ethereum/provider.py           |  2 +-
 src/ape_node/provider.py               |  2 +-
 tests/functional/geth/test_provider.py |  2 +-
 4 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 src/ape/utils/_web3_compat.py

diff --git a/src/ape/utils/_web3_compat.py b/src/ape/utils/_web3_compat.py
new file mode 100644
index 0000000000..b915e679c0
--- /dev/null
+++ b/src/ape/utils/_web3_compat.py
@@ -0,0 +1,22 @@
+from eth_account import Account as EthAccount
+
+try:
+    # Web3 v7
+    from web3.middleware import ExtraDataToPOAMiddleware  # type: ignore
+except ImportError:
+    from web3.middleware import geth_poa_middleware as ExtraDataToPOAMiddleware  # type: ignore
+
+
+def sign_hash(msghash, private_key):
+    try:
+        # Web3 v7
+        return EthAccount.unsafe_sign_hash(msghash, private_key)  # type: ignore
+    except AttributeError:
+        # Web3 v6
+        return EthAccount.signHash(msghash, private_key)  # type: ignore
+
+
+__all__ = [
+    "ExtraDataToPOAMiddleware",
+    "sign_hash",
+]
diff --git a/src/ape_ethereum/provider.py b/src/ape_ethereum/provider.py
index 95cd239977..4524bde318 100644
--- a/src/ape_ethereum/provider.py
+++ b/src/ape_ethereum/provider.py
@@ -28,7 +28,6 @@
     TransactionNotFound,
 )
 from web3.gas_strategies.rpc import rpc_gas_price_strategy
-from web3.middleware import ExtraDataToPOAMiddleware
 from web3.middleware.validation import MAX_EXTRADATA_LENGTH
 from web3.providers import AutoProvider
 from web3.providers.auto import load_provider_from_environment
@@ -57,6 +56,7 @@
 from ape.types.events import ContractLog, LogFilter
 from ape.types.gas import AutoGasLimit
 from ape.types.trace import SourceTraceback
+from ape.utils._web3_compat import ExtraDataToPOAMiddleware
 from ape.utils.basemodel import ManagerAccessMixin
 from ape.utils.misc import DEFAULT_MAX_RETRIES_TX, gas_estimation_error_message, to_int
 from ape_ethereum._print import CONSOLE_ADDRESS, console_contract
diff --git a/src/ape_node/provider.py b/src/ape_node/provider.py
index ccffadf44a..63ff0b8279 100644
--- a/src/ape_node/provider.py
+++ b/src/ape_node/provider.py
@@ -13,11 +13,11 @@
 from pydantic import field_validator
 from pydantic_settings import SettingsConfigDict
 from requests.exceptions import ConnectionError
-from web3.middleware import ExtraDataToPOAMiddleware
 
 from ape.api.config import PluginConfig
 from ape.api.providers import SubprocessProvider, TestProviderAPI
 from ape.logging import LogLevel, logger
+from ape.utils._web3_compat import ExtraDataToPOAMiddleware
 from ape.utils.misc import ZERO_ADDRESS, log_instead_of_fail, raises_not_implemented
 from ape.utils.process import JoinableQueue, spawn
 from ape.utils.testing import (
diff --git a/tests/functional/geth/test_provider.py b/tests/functional/geth/test_provider.py
index 05e0f1cc6a..a8e49e82c1 100644
--- a/tests/functional/geth/test_provider.py
+++ b/tests/functional/geth/test_provider.py
@@ -11,7 +11,6 @@
 from web3 import AutoProvider, Web3
 from web3.exceptions import ContractLogicError as Web3ContractLogicError
 from web3.exceptions import ExtraDataLengthError
-from web3.middleware import ExtraDataToPOAMiddleware
 from web3.providers import HTTPProvider
 
 from ape.exceptions import (
@@ -26,6 +25,7 @@
     VirtualMachineError,
 )
 from ape.utils import to_int
+from ape.utils._web3_compat import ExtraDataToPOAMiddleware
 from ape_ethereum.ecosystem import Block
 from ape_ethereum.provider import DEFAULT_SETTINGS, EthereumNodeProvider
 from ape_ethereum.trace import TraceApproach