simplify EntityKey type

Author: Matthias Zimmermann

README.md

@@ -119,7 +119,7 @@ SDK versions are tracked in the following files:
 Pytest is used for unit and integration testing.
 ```bash
 uv run pytest # Run all tests
-uv run pytest -k test_create_entity_simple --log-cli-level=INFO # Specific tests via keyword, print at info log level
+uv run pytest -k test_create_entity_simple --log-cli-level=info # Specific tests via keyword, print at info log level
 ```
 
 If an `.env` file is present the unit tests are run against the specifice RPC coordinates and test accounts.

src/arkiv/contract.py

@@ -9,7 +9,6 @@
 from web3.types import RPCEndpoint
 
 # Import EntityKey for munger type checking
-from .types import EntityKey
 
 STORAGE_ADDRESS: Final[ChecksumAddress] = Web3.to_checksum_address(
     "0x0000000000000000000000000000000060138453"
@@ -60,48 +59,33 @@
 ]
 
 
-def custom_munger(module: Any, *args: Any, **kwargs: Any) -> Any:
-    """Custom munger for RPC methods that automatically converts EntityKey objects."""
-    processed_args = tuple(
-        arg.value if isinstance(arg, EntityKey) else arg for arg in args
-    )
-
-    processed_kwargs = {
-        key: value.value if isinstance(value, EntityKey) else value
-        for key, value in kwargs.items()
-    }
-
-    # Apply default munger to processed arguments
-    return default_root_munger(module, *processed_args, **processed_kwargs)
-
-
 FUNCTIONS_ABI: dict[str, Method[Any]] = {
     "get_storage_value": Method(
         json_rpc_method=RPCEndpoint("golembase_getStorageValue"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "get_entity_metadata": Method(
         json_rpc_method=RPCEndpoint("golembase_getEntityMetaData"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "get_entities_to_expire_at_block": Method(
         json_rpc_method=RPCEndpoint("golembase_getEntitiesToExpireAtBlock"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "get_entity_count": Method(
         json_rpc_method=RPCEndpoint("golembase_getEntityCount"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "get_all_entity_keys": Method(
         json_rpc_method=RPCEndpoint("golembase_getAllEntityKeys"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "get_entities_of_owner": Method(
         json_rpc_method=RPCEndpoint("golembase_getEntitiesOfOwner"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
     "query_entities": Method(
         json_rpc_method=RPCEndpoint("golembase_queryEntities"),
-        mungers=[custom_munger],
+        mungers=[default_root_munger],
     ),
 }

src/arkiv/exceptions.py

@@ -1,6 +1,10 @@
 """Arkiv client exceptions."""
 
 
+class EntityKeyException(Exception):
+    pass
+
+
 class AccountNameException(Exception):
     pass
 

src/arkiv/types.py

@@ -2,89 +2,13 @@
 
 from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import NewType
 
 from eth_typing import ChecksumAddress, HexStr
 from hexbytes import HexBytes
 
-
 # Unique key for all entities
-@dataclass(frozen=True)
-class EntityKey:
-    """EntityKey dataclass that wraps HexStr for entity identification."""
-
-    value: HexStr
-
-    def __init__(self, value: str | int | HexBytes | HexStr) -> None:
-        """Create an EntityKey from various input types."""
-        if isinstance(value, str) and value.startswith("0x"):
-            # Already a hex string - validate length
-            if len(value) != 66:  # 0x + 64 hex chars
-                raise ValueError(
-                    f"EntityKey hex string must be 66 characters (0x + 64 hex), got {len(value)}"
-                )
-            object.__setattr__(self, "value", HexStr(value.lower()))
-        elif isinstance(value, int):
-            # Convert integer to hex string with 0x prefix
-            if value < 0:
-                raise ValueError("EntityKey cannot be negative")
-            hex_str = f"0x{value:064x}"  # 64 chars = 32 bytes = 256 bits
-            object.__setattr__(self, "value", HexStr(hex_str))
-        elif isinstance(value, (HexBytes, bytes)):
-            # Convert bytes to hex string
-            if len(value) != 32:  # 32 bytes = 256 bits
-                raise ValueError(
-                    f"EntityKey bytes must be exactly 32 bytes, got {len(value)}"
-                )
-            object.__setattr__(self, "value", HexStr(f"0x{value.hex()}"))
-        elif isinstance(value, str):
-            # Plain string, assume it needs 0x prefix
-            if len(value) % 2 != 0:
-                value = "0" + value  # Pad with leading zero if odd length
-            if len(value) != 64:  # Should be 64 hex characters
-                raise ValueError(
-                    f"EntityKey hex string (without 0x) must be 64 characters, got {len(value)}"
-                )
-            object.__setattr__(self, "value", HexStr(f"0x{value.lower()}"))
-        else:
-            # Try to convert via HexBytes first, then to hex string
-            try:
-                hex_bytes = HexBytes(value)
-                if len(hex_bytes) != 32:
-                    raise ValueError(
-                        f"EntityKey must represent exactly 32 bytes, got {len(hex_bytes)}"
-                    )
-                object.__setattr__(self, "value", HexStr(f"0x{hex_bytes.hex()}"))
-            except Exception as e:
-                raise ValueError(
-                    f"Cannot convert {type(value)} to EntityKey: {e}"
-                ) from e
-
-    def __str__(self) -> str:
-        """String representation."""
-        return self.value
-
-    def __repr__(self) -> str:
-        """Repr representation."""
-        return f"EntityKey('{self.value}')"
-
-    def __eq__(self, other: object) -> bool:
-        """Equality comparison."""
-        if isinstance(other, EntityKey):
-            return self.value == other.value
-        return False
-
-    def __hash__(self) -> int:
-        """Hash for use in sets/dicts."""
-        return hash(self.value)
-
-    @property
-    def hex(self) -> str:
-        """Get hex string without 0x prefix."""
-        return self.value[2:]
-
-    def to_bytes(self) -> bytes:
-        """Convert to bytes."""
-        return bytes.fromhex(self.hex)
+EntityKey = NewType("EntityKey", HexStr)
 
 
 type AnnotationValue = str | int  # Only str or non-negative int allowed

src/arkiv/utils.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 import rlp  # type: ignore[import-untyped]
+from eth_typing import HexStr
 from hexbytes import HexBytes
 from web3 import Web3
 from web3.contract import Contract
@@ -12,6 +13,7 @@
 
 from . import contract
 from .contract import STORAGE_ADDRESS
+from .exceptions import EntityKeyException
 from .types import (
     Annotation,
     AnnotationValue,
@@ -28,6 +30,52 @@
 logger = logging.getLogger(__name__)
 
 
+def to_entity_key(entity_key_int: int) -> EntityKey:
+    hex_value = Web3.to_hex(entity_key_int)
+    # ensure lenth is 66 (0x + 64 hex)
+    if len(hex_value) < 66:
+        hex_value = HexStr("0x" + hex_value[2:].zfill(64))
+    return EntityKey(hex_value)
+
+
+def entity_key_to_bytes(entity_key: EntityKey) -> bytes:
+    return bytes.fromhex(entity_key[2:])  # Strip '0x' prefix and convert to bytes
+
+
+def check_entity_key(entity_key: Any | None, label: str | None = None) -> None:
+    """Validates entity key."""
+    prefix = ""
+    if label:
+        prefix = f"{label}: "
+
+    logger.info(f"{prefix}Checking entity key {entity_key}")
+
+    if entity_key is None:
+        raise EntityKeyException("Entity key should not be None")
+    if not isinstance(entity_key, str):
+        raise EntityKeyException(
+            f"Entity key type should be str but is: {type(entity_key)}"
+        )
+    if len(entity_key) != 66:
+        raise EntityKeyException(
+            f"Entity key should be 66 characters long (0x + 64 hex) but is: {len(entity_key)}"
+        )
+    if not is_hex_str(entity_key):
+        raise EntityKeyException("Entity key should be a valid hex string")
+
+
+def is_hex_str(value: str) -> bool:
+    if not isinstance(value, str):
+        return False
+    if value.startswith("0x"):
+        value = value[2:]
+    try:
+        int(value, 16)
+        return True
+    except ValueError:
+        return False
+
+
 def to_create_operation(
     payload: bytes | None = None,
     annotations: dict[str, AnnotationValue] | None = None,
@@ -116,33 +164,36 @@ def to_receipt(
             event_args: dict[str, Any] = event_data["args"]
             event_name = event_data["event"]
 
+            entity_key: EntityKey = to_entity_key(event_args["entityKey"])
+            expiration_block: int = event_args["expirationBlock"]
+
             match event_name:
                 case contract.CREATED_EVENT:
                     creates.append(
                         CreateReceipt(
-                            entity_key=EntityKey(event_args["entityKey"]),
-                            expiration_block=int(event_args["expirationBlock"]),
+                            entity_key=entity_key,
+                            expiration_block=expiration_block,
                         )
                     )
                 case contract.UPDATED_EVENT:
                     updates.append(
                         UpdateReceipt(
-                            entity_key=EntityKey(event_args["entityKey"]),
-                            expiration_block=int(event_args["expirationBlock"]),
+                            entity_key=entity_key,
+                            expiration_block=expiration_block,
                         )
                     )
                 case contract.DELETED_EVENT:
                     deletes.append(
                         DeleteReceipt(
-                            entity_key=EntityKey(event_args["entityKey"]),
+                            entity_key=entity_key,
                         )
                     )
                 case contract.EXTENDED_EVENT:
                     extensions.append(
                         ExtendReceipt(
-                            entity_key=EntityKey(event_args["entityKey"]),
-                            old_expiration_block=int(event_args["oldExpirationBlock"]),
-                            new_expiration_block=int(event_args["newExpirationBlock"]),
+                            entity_key=entity_key,
+                            old_expiration_block=event_args["oldExpirationBlock"],
+                            new_expiration_block=event_args["newExpirationBlock"],
                         )
                     )
                 case _:
@@ -195,7 +246,7 @@ def format_annotation(annotation: Annotation) -> tuple[str, AnnotationValue]:
         # Update
         [
             [
-                element.entity_key.to_bytes(),
+                entity_key_to_bytes(element.entity_key),
                 element.btl,
                 element.data,
                 list(map(format_annotation, element.string_annotations)),
@@ -206,14 +257,14 @@ def format_annotation(annotation: Annotation) -> tuple[str, AnnotationValue]:
         # Delete
         [
             [
-                element.entity_key.to_bytes(),
+                entity_key_to_bytes(element.entity_key),
             ]
             for element in tx.deletes
         ],
         # Extend
         [
             [
-                element.entity_key.to_bytes(),
+                entity_key_to_bytes(element.entity_key),
                 element.number_of_blocks,
             ]
             for element in tx.extensions

tests/test_entity_create.py

@@ -7,24 +7,14 @@
 
 from arkiv.client import Arkiv
 from arkiv.contract import STORAGE_ADDRESS
-from arkiv.types import EntityKey, Operations
-from arkiv.utils import to_create_operation, to_receipt, to_tx_params
+from arkiv.types import Operations
+from arkiv.utils import check_entity_key, to_create_operation, to_receipt, to_tx_params
 
 logger = logging.getLogger(__name__)
 
 TX_SUCCESS = 1
 
 
-def check_entity_key(label: str, entity_key: EntityKey) -> None:
-    """Check entity key validity."""
-    logger.info(f"{label}: Checking entity key {entity_key}")
-    assert entity_key is not None, f"{label}: Entity key should not be None"
-    assert isinstance(entity_key, EntityKey), f"{label}: Entity key should be EntityKey"
-    assert len(entity_key.to_bytes()) == 32, (
-        f"{label}: Entity key should be 32 bytes long"
-    )
-
-
 def check_tx_hash(label: str, tx_hash: HexBytes) -> None:
     """Check transaction hash validity."""
     logger.info(f"{label}: Checking transaction hash {tx_hash.to_0x_hex()}")
@@ -161,7 +151,7 @@ def test_create_entity_simple(self, arkiv_client_http: Arkiv) -> None:
         )
 
         label = "create_entity (a)"
-        check_entity_key(label, entity_key)
+        check_entity_key(entity_key, label)
         check_tx_hash(label, tx_hash)
 
         entity = arkiv_client_http.arkiv.get_entity(entity_key)

tests/test_entity_create_parallel.py

@@ -115,9 +115,8 @@ def client_task(client: Arkiv, client_idx: int, num_entities: int) -> list[Entit
 @pytest.mark.parametrize(
     "num_clients,num_entities",
     [
-        (2, 10),
-        # (2, 3),
-        # (4, 2),
+        # (20, 50),
+        (2, 3),
     ],
 )
 def test_parallel_entity_creation(
@@ -130,6 +129,7 @@ def test_parallel_entity_creation(
     if not rpc_url:
         pytest.skip("No Arkiv node available for testing")
 
+    # Create Arkiv clients
     logger.info(f"Starting {num_clients} Arkiv clients...")
     clients = []
     for i in range(num_clients):
@@ -138,6 +138,8 @@ def test_parallel_entity_creation(
         client = create_client(container, rpc_url, client_idx)
         account: ChecksumAddress = cast(ChecksumAddress, client.eth.default_account)
         balance = client.eth.get_balance(account)
+
+        # Only use clients with non-zero balance
         if balance > 0:
             logger.info(f"Arkiv client[{client_idx}] started.")
             clients.append(client)
@@ -147,6 +149,7 @@ def test_parallel_entity_creation(
     # Remember start time
     start_time = time.time()
 
+    # Start all clients in separate threads (pseudo-parallelism)
     threads = []
     for client_idx in range(len(clients)):
         client = clients[client_idx]
@@ -164,6 +167,7 @@ def test_parallel_entity_creation(
     end_time = time.time()
     elapsed_time = end_time - start_time
 
+    logger.info(f"Total active clients: {len(clients)}")
     logger.info(f"Total successful entity creation TX: {tx_counter}")
     logger.info(f"TX creation start time: {start_time:.2f}, end time: {end_time:.2f}")
     logger.info(f"All Arkiv clients have completed in {elapsed_time:.2f} seconds.")