types/storage: add auto store counter

src/ethereum_test_tools/common/types.py

@@ -4,7 +4,20 @@
 import json
 from copy import copy, deepcopy
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeAlias
+from itertools import count
+from typing import (
+    Any,
+    ClassVar,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeAlias,
+)
 
 from coincurve.keys import PrivateKey, PublicKey
 from ethereum import rlp as eth_rlp
@@ -84,6 +97,8 @@ class Storage:
 
     data: Dict[int, int]
 
+    current_slot: Iterator[int]
+
     StorageDictType: ClassVar[TypeAlias] = Dict[str | int | bytes, str | int | bytes]
     """
     Dictionary type to be used when defining an input to initialize a storage.
@@ -225,7 +240,7 @@ def key_value_to_string(value: int) -> str:
             hex_str = "0" + hex_str
         return "0x" + hex_str
 
-    def __init__(self, input: StorageDictType):
+    def __init__(self, input: StorageDictType, start_slot: int = 0):
         """
         Initializes the storage using a given mapping which can have
         keys and values either as string or int.
@@ -237,7 +252,7 @@ def __init__(self, input: StorageDictType):
             value = Storage.parse_key_value(input[key])
             key = Storage.parse_key_value(key)
             self.data[key] = value
-        pass
+        self.current_slot = count(start_slot)
 
     def __len__(self) -> int:
         """Returns number of elements in the storage"""
@@ -263,6 +278,16 @@ def __delitem__(self, key: str | int):
         """Deletes an item from the storage"""
         del self.data[Storage.parse_key_value(key)]
 
+    def store_next(self, value: str | int) -> int:
+        """
+        Stores a value in the storage and returns the key where the value is stored.
+
+        Increments the key counter so the next time this function is called,
+        the next key is used.
+        """
+        self[slot := next(self.current_slot)] = value
+        return slot
+
     def to_dict(self) -> Mapping[str, str]:
         """
         Converts the storage into a string dict with appropriate 32-byte

src/ethereum_test_tools/tests/test_types.py

@@ -65,6 +65,18 @@ def test_storage():
     with pytest.raises(Storage.AmbiguousKeyValue):
         s.to_dict()
 
+    # Check store counter
+    s = Storage({})
+    s.store_next(0x100)
+    s.store_next(0x200)
+    s.store_next(0x300)
+    d = s.to_dict()
+    assert d == {
+        "0x00": ("0x0100"),
+        "0x01": ("0x0200"),
+        "0x02": ("0x0300"),
+    }
+
 
 @pytest.mark.parametrize(
     ["account", "alloc", "should_pass"],