From 88f8255e3f876da38e67ba4e9833658101a6ae22 Mon Sep 17 00:00:00 2001
From: replikeit <egorprihodko8@gmail.com>
Date: Tue, 22 Oct 2024 16:29:53 +0300
Subject: [PATCH 1/2] Update: add feature

---
 pyrevm.pyi      |  8 ++++++++
 src/database.rs | 14 ++++++++++++++
 src/evm.rs      | 15 ++++++++++++---
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/pyrevm.pyi b/pyrevm.pyi
index 5102686..4700a7e 100644
--- a/pyrevm.pyi
+++ b/pyrevm.pyi
@@ -201,6 +201,14 @@ class EVM:
         :param info: The account info.
         """
 
+    def insert_account_storage(self: "EVM", address: str, index: int, value: int) -> None:
+        """
+        Inserts the provided value for slot of in the database at the specified address
+        :param address: The address of the account.
+        :param index: slot in storage.
+        :param value: value for slot.
+        """
+
     def message_call(
         self: "EVM",
         caller: str,
diff --git a/src/database.rs b/src/database.rs
index 3022354..9d8a80a 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -47,6 +47,20 @@ impl DB {
         }
     }
 
+    pub(crate) fn insert_insert_account_storage(
+        &mut self,
+        address: Address,
+        slot: U256,
+        value: U256,
+    ) -> PyResult<()> {
+        match self {
+            DB::Memory(db) =>
+                db.insert_account_storage(address, slot, value).map_err(pyerr),
+            DB::Fork(db) =>
+                db.insert_account_storage(address, slot, value).map_err(pyerr),
+        }
+    }
+
     pub(crate) fn get_accounts(&self) -> &HashMap<Address, DbAccount> {
         match self {
             DB::Memory(db) => &db.accounts,
diff --git a/src/evm.rs b/src/evm.rs
index a5d835f..90fed84 100644
--- a/src/evm.rs
+++ b/src/evm.rs
@@ -52,7 +52,8 @@ pub struct EVM {
 impl EVM {
     /// Create a new EVM instance.
     #[new]
-    #[pyo3(signature = (env = None, fork_url = None, fork_block = None, gas_limit = 18446744073709551615, tracing = false, spec_id = "LATEST"))]
+    #[pyo3(signature = (env = None, fork_url = None, fork_block = None, gas_limit = 18446744073709551615, tracing = false, spec_id = "LATEST")
+    )]
     fn new(
         env: Option<Env>,
         fork_url: Option<&str>,
@@ -151,6 +152,12 @@ impl EVM {
         Ok(())
     }
 
+    /// Inserts the provided value for slot of in the database at the specified address
+    fn insert_account_storage(&mut self, address: &str, index: U256, value: U256) -> PyResult<()> {
+        let target = addr(address)?;
+        self.context.db.insert_insert_account_storage(target, index, value)
+    }
+
     /// Set the balance of a given address.
     fn set_balance(&mut self, address: &str, balance: U256) -> PyResult<()> {
         let address = addr(address)?;
@@ -166,7 +173,8 @@ impl EVM {
         Ok(balance)
     }
 
-    #[pyo3(signature = (caller, to, calldata = None, value = None, gas = None, gas_price = None, is_static = false))]
+    #[pyo3(signature = (caller, to, calldata = None, value = None, gas = None, gas_price = None, is_static = false)
+    )]
     pub fn message_call(
         &mut self,
         caller: &str,
@@ -193,7 +201,8 @@ impl EVM {
     }
 
     /// Deploy a contract with the given code.
-    #[pyo3(signature = (deployer, code, value = None, gas = None, gas_price = None, is_static = false, _abi = None))]
+    #[pyo3(signature = (deployer, code, value = None, gas = None, gas_price = None, is_static = false, _abi = None)
+    )]
     fn deploy(
         &mut self,
         deployer: &str,

From 83ca64cbbb3f8d2b4b8d53b43ed6e942b3902f00 Mon Sep 17 00:00:00 2001
From: replikeit <egorprihodko8@gmail.com>
Date: Tue, 22 Oct 2024 17:22:20 +0300
Subject: [PATCH 2/2] Update: add tests

---
 Cargo.lock        |  2 +-
 Cargo.toml        |  2 +-
 tests/test_evm.py | 14 +++++++++++---
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 434e11c..c10c00b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1817,7 +1817,7 @@ dependencies = [
 
 [[package]]
 name = "pyrevm"
-version = "0.3.3"
+version = "0.3.4"
 dependencies = [
  "ethers-core",
  "ethers-providers",
diff --git a/Cargo.toml b/Cargo.toml
index 61e8514..a397903 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pyrevm"
-version = "0.3.3"
+version = "0.3.4"
 edition = "2021"
 
 [lib]
diff --git a/tests/test_evm.py b/tests/test_evm.py
index ecf0594..4ce25e1 100644
--- a/tests/test_evm.py
+++ b/tests/test_evm.py
@@ -11,8 +11,8 @@
 
 # use your own key during development to avoid rate limiting the CI job
 fork_url = (
-    os.getenv("FORK_URL")
-    or "https://mainnet.infura.io/v3/c60b0bb42f8a4c6481ecd229eddaca27"
+        os.getenv("FORK_URL")
+        or "https://mainnet.infura.io/v3/c60b0bb42f8a4c6481ecd229eddaca27"
 )
 
 KWARG_CASES = [
@@ -60,6 +60,14 @@ def test_fork_storage():
     assert value > 0
 
 
+def test_set_into_storage():
+    weth = "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"
+    evm = EVM(fork_url=fork_url, fork_block="latest")
+    evm.insert_account_storage(weth, 0, 10)
+    value = evm.storage(weth, 0)
+    assert value == 10
+
+
 def test_deploy():
     evm = EVM()
 
@@ -265,7 +273,7 @@ def test_tx_setters():
 
 @pytest.mark.parametrize(
     "excess_blob_gas,expected_fee",
-    [(0, 1), (10**3, 1), (2**24, 152), (2**26, 537070730)],
+    [(0, 1), (10 ** 3, 1), (2 ** 24, 152), (2 ** 26, 537070730)],
 )
 def test_get_blobbasefee(excess_blob_gas, expected_fee):
     evm = EVM()