add support for hardhat setStorageAt/setCode/setBalance

boa/network.py

@@ -610,11 +610,35 @@ def get_chain_id(self) -> int:
         chain_id = self._rpc.fetch("eth_chainId", [])
         return int(chain_id, 16)
 
+    @contextlib.contextmanager
+    def sender(self, address):
+        self._rpc.fetch("hardhat_impersonateAccount", [to_hex(addr)])
+        tmp = self.eoa
+        addr = Address(address)
+        self.add_account(ExternalAccount(addr, self._rpc), force_eoa=True)
+        try:
+            yield
+        finally:
+            self.eoa = tmp
+
     def set_balance(self, address, value):
-        raise NotImplementedError("Cannot use set_balance in network mode")
+        super().set_balance(address, value)  # set it in the fork
+
+        self._rpc.fetch(
+            "hardhat_setBalance", [to_hex(address), to_hex(value, pad_nibbles=64)]
+        )
 
     def set_code(self, address: _AddressType, code: bytes) -> None:
-        raise NotImplementedError("Cannot use set_code in network mode")
+        super().set_code(address, code)  # set it in the fork
+
+        self._rpc.fetch("hardhat_setCode", [to_hex(address), to_hex(code)])
 
     def set_storage(self, address: _AddressType, slot: int, value: int) -> None:
-        raise NotImplementedError("Cannot use set_storage in network mode")
+        super().set_storage(address, slot, value)  # set it in the fork
+
+        # will throw if the provider does not have hardhat_setStorageAt
+        # hardhat requires padded hex values
+        self._rpc.fetch(
+            "hardhat_setStorageAt",
+            [to_hex(address), to_hex(slot), to_hex(value, pad_nibbles=64)],
+        )

boa/rpc.py

@@ -23,14 +23,18 @@ def fixup_dict(kv):
     return {k: to_hex(v) for (k, v) in trim_dict(kv).items()}
 
 
-def to_hex(s: int | bytes | str) -> str:
+def pad_hex(s: str, pad_nibbles: int) -> str:
+    return "0x" + s.removeprefix("0x").rjust(pad_nibbles, "0")
+
+
+def to_hex(s: int | bytes | str, pad_nibbles=0) -> str:
     if isinstance(s, int):
-        return hex(s)
+        return pad_hex(hex(s), pad_nibbles)
     if isinstance(s, bytes):
-        return "0x" + s.hex()
+        return pad_hex("0x" + s.hex(), pad_nibbles)
     if isinstance(s, str):
         assert s.startswith("0x")
-        return s
+        return pad_hex(s, pad_nibbles)
     raise TypeError(
         f"to_hex expects bytes, int or (hex) string, but got {type(s)}: {s}"
     )