Merge pull request #1 from p2p-org/add_insert_storage

Add insert storage

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyrevm"
-version = "0.3.3"
+version = "0.3.4"
 edition = "2021"
 
 [lib]

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,

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,

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,

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()