simplify Annotations, refactorings

Author: Matthias Zimmermann

src/arkiv/account.py

@@ -175,6 +175,10 @@ def main() -> None:
     password = getpass.getpass("Enter wallet password: ")
     encrypted = account.local_account.encrypt(password)
 
+    # Ensure address has 0x prefix (eth_account.encrypt doesn't include it)
+    if "address" in encrypted and not encrypted["address"].startswith("0x"):
+        encrypted["address"] = "0x" + encrypted["address"]
+
     with wallet_path.open("w") as f:
         json.dump(encrypted, f)
 

src/arkiv/exceptions.py

@@ -5,6 +5,10 @@ class EntityKeyException(Exception):
     pass
 
 
+class AnnotationException(Exception):
+    pass
+
+
 class AccountNameException(Exception):
     pass
 

src/arkiv/module.py

@@ -4,32 +4,32 @@
 import logging
 from typing import TYPE_CHECKING, Any
 
+from eth_typing import ChecksumAddress
 from hexbytes import HexBytes
 from web3 import Web3
 from web3.types import TxParams, TxReceipt
 
 from .contract import EVENTS_ABI, FUNCTIONS_ABI, STORAGE_ADDRESS
 from .types import (
-    AnnotationValue,
+    ALL,
+    ANNOTATIONS,
+    METADATA,
+    PAYLOAD,
+    Annotations,
+    CreateOp,
     Entity,
     EntityKey,
-    Metadata,
     Operations,
     TransactionReceipt,
 )
-from .utils import merge_annotations, to_create_operation, to_receipt, to_tx_params
+from .utils import merge_annotations, to_receipt, to_tx_params
 
 # Deal with potential circular imports between client.py and module.py
 if TYPE_CHECKING:
     from .client import Arkiv
 
 logger = logging.getLogger(__name__)
 
-PAYLOAD = 1
-ANNOTATIONS = 2
-METADATA = 4
-ALL = PAYLOAD | ANNOTATIONS | METADATA
-
 TX_SUCCESS = 1
 
 
@@ -57,7 +57,7 @@ def is_available(self) -> bool:
     def create_entity(
         self,
         payload: bytes | None = None,
-        annotations: dict[str, AnnotationValue] | None = None,
+        annotations: Annotations | None = None,
         btl: int = 0,
         tx_params: TxParams | None = None,
     ) -> tuple[EntityKey, HexBytes]:
@@ -73,10 +73,14 @@ def create_entity(
         Returns:
             Transaction hash of the create operation
         """
+        # Check and set defaults
+        if not payload:
+            payload = b""
+        if not annotations:
+            annotations = Annotations({})
+
         # Create the operation
-        create_op = to_create_operation(
-            payload=payload, annotations=annotations, btl=btl
-        )
+        create_op = CreateOp(payload=payload, annotations=annotations, btl=btl)
 
         # Wrap in Operations container
         operations = Operations(creates=[create_op])
@@ -118,9 +122,10 @@ def get_entity(self, entity_key: EntityKey, fields: int = ALL) -> Entity:
             Entity object with the requested fields
         """
         # Gather the requested data
-        payload = None
-        metadata = None
-        annotations = None
+        owner: ChecksumAddress | None = None
+        expires_at_block: int | None = None
+        payload: bytes | None = None
+        annotations: Annotations | None = None
 
         # HINT: rpc methods to fetch entity content might change this is the place to adapt
         # get and decode payload if requested
@@ -133,21 +138,11 @@ def get_entity(self, entity_key: EntityKey, fields: int = ALL) -> Entity:
 
             if fields & METADATA:
                 # Convert owner address to checksummed format
-                owner_address = metadata_all.get("owner")
-                if not owner_address:
-                    raise ValueError("Entity metadata missing required owner field")
-                checksummed_owner = Web3.to_checksum_address(owner_address)
-
-                expires_at_block = metadata_all.get("expiresAtBlock")
-                if expires_at_block is None:
-                    raise ValueError(
-                        "Entity metadata missing required expiresAtBlock field"
-                    )
-
-                metadata = Metadata(
-                    owner=checksummed_owner,
-                    expires_at_block=int(expires_at_block),
-                )
+                owner_metadata = metadata_all.get("owner")
+                owner = Web3.to_checksum_address(owner_metadata)
+
+                expires_at_block_metadata = metadata_all.get("expiresAtBlock")
+                expires_at_block = int(expires_at_block_metadata)
 
             if fields & ANNOTATIONS:
                 annotations = merge_annotations(
@@ -158,8 +153,10 @@ def get_entity(self, entity_key: EntityKey, fields: int = ALL) -> Entity:
         # Create and return entity
         return Entity(
             entity_key=entity_key,
+            fields=fields,
+            owner=owner,
+            expires_at_block=expires_at_block,
             payload=payload,
-            metadata=metadata,
             annotations=annotations,
         )
 

src/arkiv/types.py

@@ -6,71 +6,57 @@
 
 from eth_typing import ChecksumAddress, HexStr
 from hexbytes import HexBytes
+from web3.datastructures import AttributeDict
 
-# Unique key for all entities
-EntityKey = NewType("EntityKey", HexStr)
-
-
-type AnnotationValue = str | int  # Only str or non-negative int allowed
-
-
-@dataclass(frozen=True)
-class Annotation:
-    """Class to represent annotations with string or non-negative integer values."""
-
-    key: str
-    value: AnnotationValue
-
-    def __post_init__(self) -> None:
-        """Validate that integer values are non-negative."""
-        if isinstance(self.value, int) and self.value < 0:
-            raise ValueError(
-                f"Integer annotation values must be non-negative, got: {self.value}"
-            )
-
-    # @override
-    def __repr__(self) -> str:
-        """Encode annotation as a string."""
-        return f"{type(self).__name__}({self.key} -> {self.value})"
+# Field bitmask values to specify which entity fields are populated
+PAYLOAD = 1
+ANNOTATIONS = 2
+METADATA = 4
+ALL = PAYLOAD | ANNOTATIONS | METADATA
 
 
-@dataclass(frozen=True)
-class Metadata:
-    """A class representing entity metadata."""
+# Unique key for all entities
+EntityKey = NewType("EntityKey", HexStr)
 
-    owner: ChecksumAddress
-    expires_at_block: int
+# Entity annotations
+Annotations = NewType("Annotations", dict[str, str | int])
 
 
 @dataclass(frozen=True)
 class Entity:
     """A class representing an entity."""
 
-    entity_key: EntityKey
-    metadata: Metadata | None
+    entity_key: EntityKey  # Unique identifier for the entity
+    fields: int  # Bitmask representing which fields are populated
+
+    # Populated when fields | METADATA returns true
+    owner: ChecksumAddress | None
+    expires_at_block: int | None
+
+    # Populated when fields | PAYLOAD returns true
     payload: bytes | None
-    annotations: dict[str, AnnotationValue] | None
+
+    # Populated when fields | ANNOTATIONS returns true
+    annotations: Annotations | None
 
 
 @dataclass(frozen=True)
 class CreateOp:
     """Class to represent a create operation."""
 
-    data: bytes
+    payload: bytes
+    annotations: Annotations
     btl: int
-    string_annotations: Sequence[Annotation]
-    numeric_annotations: Sequence[Annotation]
 
 
 @dataclass(frozen=True)
 class UpdateOp:
     """Class to represent an update operation."""
 
     entity_key: EntityKey
-    data: bytes
+    payload: bytes
+    annotations: Annotations
     btl: int
-    string_annotations: Sequence[Annotation]
-    numeric_annotations: Sequence[Annotation]
 
 
 @dataclass(frozen=True)
@@ -164,3 +150,12 @@ class TransactionReceipt:
     updates: Sequence[UpdateReceipt]
     extensions: Sequence[ExtendReceipt]
     deletes: Sequence[DeleteReceipt]
+
+
+# Low level annotations for RLP encoding
+StringAnnotationsRlp = NewType("StringAnnotationsRlp", list[tuple[str, str]])
+NumericAnnotationsRlp = NewType("NumericAnnotationsRlp", list[tuple[str, int]])
+
+# Low level annotations for entity decoding
+StringAnnotations = NewType("StringAnnotations", AttributeDict[str, str])
+NumericAnnotations = NewType("NumericAnnotations", AttributeDict[str, int])