refactor entity ops watching

Author: matthiaszimmermann

src/arkiv/events.py

@@ -11,11 +11,13 @@
 from web3.contract.contract import ContractEvent
 from web3.types import EventData, LogReceipt
 
-from .contract import CREATED_EVENT, UPDATED_EVENT
+from .contract import CREATED_EVENT, EXTENDED_EVENT, UPDATED_EVENT
 from .types import (
     CreateCallback,
     CreateEvent,
     EventType,
+    ExtendCallback,
+    ExtendEvent,
     TxHash,
     UpdateCallback,
     UpdateEvent,
@@ -32,7 +34,7 @@ def __init__(
         self,
         contract: Contract,
         event_type: EventType,
-        callback: CreateCallback | UpdateCallback,
+        callback: CreateCallback | UpdateCallback | ExtendCallback,
         from_block: str | int = "latest",
         auto_start: bool = True,
     ) -> None:
@@ -48,7 +50,7 @@ def __init__(
         """
         self.contract: Contract = contract
         self.event_type: EventType = event_type
-        self.callback: CreateCallback | UpdateCallback = callback
+        self.callback: CreateCallback | UpdateCallback | ExtendCallback = callback
         self.from_block: str | int = from_block
 
         # Internal state
@@ -78,6 +80,9 @@ def start(self) -> None:
         elif self.event_type == "updated":
             contract_event = self.contract.events[UPDATED_EVENT]
             self._filter = contract_event.create_filter(from_block=self.from_block)
+        elif self.event_type == "extended":
+            contract_event = self.contract.events[EXTENDED_EVENT]
+            self._filter = contract_event.create_filter(from_block=self.from_block)
         else:
             raise NotImplementedError(
                 f"Event type {self.event_type} not yet implemented"
@@ -166,50 +171,72 @@ def _process_event(self, event_data: EventData) -> None:
         """
         logger.info(f"Processing event: {event_data}")
 
-        # Extract data based on event type
-        if self.event_type == "created":
-            # Parse CreateEvent
-            entity_key = to_entity_key(event_data["args"]["entityKey"])
-            expiration_block = event_data["args"]["expirationBlock"]
-
-            # Ensure transaction hash has 0x prefix
-            tx_hash_hex = event_data["transactionHash"].hex()
-            if not tx_hash_hex.startswith("0x"):
-                tx_hash_hex = f"0x{tx_hash_hex}"
-            tx_hash = TxHash(HexStr(tx_hash_hex))
+        # Extract common data
+        entity_key = to_entity_key(event_data["args"]["entityKey"])
+        tx_hash = self._extract_tx_hash(event_data)
 
+        # Create event object and trigger callback based on type
+        if self.event_type == "created":
             create_event = CreateEvent(
-                entity_key=entity_key, expiration_block=expiration_block
+                entity_key=entity_key,
+                expiration_block=event_data["args"]["expirationBlock"],
+            )
+            self._trigger_callback(
+                cast(CreateCallback, self.callback), create_event, tx_hash
             )
-
-            # Trigger callback
-            try:
-                # Type narrowing: when event_type is "created", callback is CreateCallback
-                cast(CreateCallback, self.callback)(create_event, tx_hash)
-            except Exception as e:
-                logger.error(f"Error in callback: {e}", exc_info=True)
 
         elif self.event_type == "updated":
-            # Parse UpdateEvent
-            entity_key = to_entity_key(event_data["args"]["entityKey"])
-            expiration_block = event_data["args"]["expirationBlock"]
-
-            # Ensure transaction hash has 0x prefix
-            tx_hash_hex = event_data["transactionHash"].hex()
-            if not tx_hash_hex.startswith("0x"):
-                tx_hash_hex = f"0x{tx_hash_hex}"
-            tx_hash = TxHash(HexStr(tx_hash_hex))
-
             update_event = UpdateEvent(
-                entity_key=entity_key, expiration_block=expiration_block
+                entity_key=entity_key,
+                expiration_block=event_data["args"]["expirationBlock"],
+            )
+            self._trigger_callback(
+                cast(UpdateCallback, self.callback), update_event, tx_hash
             )
 
-            # Trigger callback
-            try:
-                # Type narrowing: when event_type is "updated", callback is UpdateCallback
-                cast(UpdateCallback, self.callback)(update_event, tx_hash)
-            except Exception as e:
-                logger.error(f"Error in callback: {e}", exc_info=True)
+        elif self.event_type == "extended":
+            extend_event = ExtendEvent(
+                entity_key=entity_key,
+                old_expiration_block=event_data["args"]["oldExpirationBlock"],
+                new_expiration_block=event_data["args"]["newExpirationBlock"],
+            )
+            self._trigger_callback(
+                cast(ExtendCallback, self.callback), extend_event, tx_hash
+            )
 
         else:
             logger.warning(f"Unknown event type: {self.event_type}")
+
+    def _extract_tx_hash(self, event_data: EventData) -> TxHash:
+        """
+        Extract and normalize transaction hash from event data.
+
+        Args:
+            event_data: Event data from Web3 filter
+
+        Returns:
+            Transaction hash with 0x prefix
+        """
+        tx_hash_hex = event_data["transactionHash"].hex()
+        if not tx_hash_hex.startswith("0x"):
+            tx_hash_hex = f"0x{tx_hash_hex}"
+        return TxHash(HexStr(tx_hash_hex))
+
+    def _trigger_callback(
+        self,
+        callback: CreateCallback | UpdateCallback | ExtendCallback,
+        event: CreateEvent | UpdateEvent | ExtendEvent,
+        tx_hash: TxHash,
+    ) -> None:
+        """
+        Trigger callback with error handling.
+
+        Args:
+            callback: Callback function to invoke
+            event: Event object to pass to callback
+            tx_hash: Transaction hash to pass to callback
+        """
+        try:
+            callback(event, tx_hash)  # type: ignore[arg-type]
+        except Exception as e:
+            logger.error(f"Error in callback: {e}", exc_info=True)

src/arkiv/module.py

@@ -23,6 +23,8 @@
     DeleteOp,
     Entity,
     EntityKey,
+    EventType,
+    ExtendCallback,
     ExtendOp,
     Operations,
     TransactionReceipt,
@@ -397,35 +399,16 @@ def watch_entity_created(
         """
         Watch for entity creation events.
 
-        Args:
-            callback: Function to call when a creation event is detected.
-                     Receives (CreateEvent, TxHash) as arguments.
-            from_block: Starting block for the filter ('latest' or block number)
-            auto_start: If True, starts polling immediately
-
-        Returns:
-            EventFilter instance for controlling the watch
+        Creates an event filter that monitors entity creation events. The callback
+        receives (CreateEvent, TxHash) for each created entity.
 
-        Example:
-            def on_create(event: CreateEvent, tx_hash: TxHash) -> None:
-                print(f"Entity created: {event.entity_key}")
-
-            filter = arkiv.watch_entity_created(on_create)
-            # ... later ...
-            filter.stop()
+        See `_watch_entity_event` for detailed documentation on parameters, return
+        value, error handling, and usage examples.
         """
-        event_filter = EventFilter(
-            contract=self.contract,
-            event_type="created",
-            callback=callback,
-            from_block=from_block,
-            auto_start=auto_start,
+        return self._watch_entity_event(
+            "created", callback, from_block=from_block, auto_start=auto_start
         )
 
-        # Track the filter for cleanup
-        self._active_filters.append(event_filter)
-        return event_filter
-
     def watch_entity_updated(
         self,
         callback: UpdateCallback,
@@ -436,15 +419,87 @@ def watch_entity_updated(
         """
         Watch for entity update events.
 
-        This method creates an event filter that monitors for entity updates on the
-        Arkiv storage contract. The callback is invoked each time an entity is updated,
-        receiving details about the update event and the transaction hash.
+        Creates an event filter that monitors entity update events. The callback
+        receives (UpdateEvent, TxHash) for each updated entity.
+
+        See `_watch_entity_event` for detailed documentation on parameters, return
+        value, error handling, and usage examples.
+        """
+        return self._watch_entity_event(
+            "updated", callback, from_block=from_block, auto_start=auto_start
+        )
+
+    def watch_entity_extended(
+        self,
+        callback: ExtendCallback,
+        *,
+        from_block: str | int = "latest",
+        auto_start: bool = True,
+    ) -> EventFilter:
+        """
+        Watch for entity extension events.
+
+        Creates an event filter that monitors entity lifetime extension events. The
+        callback receives (ExtendEvent, TxHash) for each extended entity.
+
+        See `_watch_entity_event` for detailed documentation on parameters, return
+        value, error handling, and usage examples.
+        """
+        return self._watch_entity_event(
+            "extended", callback, from_block=from_block, auto_start=auto_start
+        )
+
+    def cleanup_filters(self) -> None:
+        """
+        Stop and uninstall all active event filters.
+
+        This is automatically called when the Arkiv client exits its context,
+        but can be called manually if needed.
+        """
+        if not self._active_filters:
+            logger.debug("No active filters to cleanup")
+            return
+
+        logger.info(
+            f"Cleaning up {len(self._active_filters)} active event filter(s)..."
+        )
+
+        for event_filter in self._active_filters:
+            try:
+                event_filter.uninstall()
+            except Exception as e:
+                logger.warning(f"Error cleaning up filter: {e}")
+
+        self._active_filters.clear()
+        logger.info("All event filters cleaned up")
+
+    @property
+    def active_filters(self) -> list[EventFilter]:
+        """Get a copy of currently active event filters."""
+        return list(self._active_filters)
+
+    def _watch_entity_event(
+        self,
+        event_type: EventType,
+        callback: CreateCallback | UpdateCallback | ExtendCallback,
+        *,
+        from_block: str | int = "latest",
+        auto_start: bool = True,
+    ) -> EventFilter:
+        """
+        Internal method to watch for entity events.
+
+        This method creates an event filter that monitors for entity events on the
+        Arkiv storage contract. The callback is invoked each time the specified event
+        occurs, receiving details about the event and the transaction hash.
 
         Args:
-            callback: Function to call when an update event is detected.
-                     Receives (UpdateEvent, TxHash) as arguments.
+            event_type: Type of event to watch for ("created", "updated", "extended")
+            callback: Function to call when an event is detected.
+                     Receives (Event, TxHash) as arguments where Event is one of:
+                     CreateEvent, UpdateEvent, or ExtendEvent depending on event_type.
             from_block: Starting block for the filter. Can be:
-                       - "latest": Only watch for new updates (default)
+                       - "latest": Only watch for new events (default)
                        - Block number (int): Watch from a specific historical block
             auto_start: If True, starts polling immediately (default: True).
                        If False, you must manually call filter.start()
@@ -462,44 +517,44 @@ def watch_entity_updated(
 
         Example:
             Basic usage with automatic start:
-                >>> def on_update(event: UpdateEvent, tx_hash: TxHash) -> None:
-                ...     print(f"Entity updated: {event.entity_key}")
-                ...     print(f"New expiration: {event.expiration_block}")
+                >>> def on_event(event: CreateEvent, tx_hash: TxHash) -> None:
+                ...     print(f"Event occurred: {event.entity_key}")
                 ...
-                >>> filter = arkiv.watch_entity_updated(on_update)
-                >>> # Filter is now running and will call on_update for each update
+                >>> filter = arkiv._watch_entity_event("created", on_event)
+                >>> # Filter is now running and will call on_event for each event
                 >>> # ... later ...
                 >>> filter.stop()  # Pause watching
                 >>> filter.uninstall()  # Cleanup resources
 
             Manual start/stop control:
-                >>> def on_update(event: UpdateEvent, tx_hash: TxHash) -> None:
-                ...     print(f"Updated: {event.entity_key}")
+                >>> def on_event(event: UpdateEvent, tx_hash: TxHash) -> None:
+                ...     print(f"Event occurred: {event.entity_key}")
                 ...
-                >>> filter = arkiv.watch_entity_updated(on_update, auto_start=False)
+                >>> filter = arkiv._watch_entity_event("updated", on_event, auto_start=False)
                 >>> # Do some setup work...
                 >>> filter.start()  # Begin watching
                 >>> # ... later ...
                 >>> filter.stop()  # Stop watching
                 >>> filter.uninstall()  # Cleanup
 
-            Historical updates from specific block:
-                >>> filter = arkiv.watch_entity_updated(
-                ...     on_update,
+            Historical events from specific block:
+                >>> filter = arkiv._watch_entity_event(
+                ...     "extended",
+                ...     on_event,
                 ...     from_block=1000  # Start from block 1000
                 ... )
 
         Note:
-            - Only captures UPDATE events (not creates, deletes, or extends)
-            - With from_block="latest", misses updates before filter creation
+            - Only captures the specified event type (not other lifecycle events)
+            - With from_block="latest", misses events before filter creation
             - Filter must be uninstalled via filter.uninstall() to free resources
             - All active filters are automatically cleaned up when Arkiv client
               context exits
             - Callback exceptions are caught and logged but don't stop the filter
         """
         event_filter = EventFilter(
             contract=self.contract,
-            event_type="updated",
+            event_type=event_type,
             callback=callback,
             from_block=from_block,
             auto_start=auto_start,
@@ -509,35 +564,6 @@ def watch_entity_updated(
         self._active_filters.append(event_filter)
         return event_filter
 
-    @property
-    def active_filters(self) -> list[EventFilter]:
-        """Get a copy of currently active event filters."""
-        return list(self._active_filters)
-
-    def cleanup_filters(self) -> None:
-        """
-        Stop and uninstall all active event filters.
-
-        This is automatically called when the Arkiv client exits its context,
-        but can be called manually if needed.
-        """
-        if not self._active_filters:
-            logger.debug("No active filters to cleanup")
-            return
-
-        logger.info(
-            f"Cleaning up {len(self._active_filters)} active event filter(s)..."
-        )
-
-        for event_filter in self._active_filters:
-            try:
-                event_filter.uninstall()
-            except Exception as e:
-                logger.warning(f"Error cleaning up filter: {e}")
-
-        self._active_filters.clear()
-        logger.info("All event filters cleaned up")
-
     def _get_owner(self, metadata: dict[str, Any]) -> ChecksumAddress:
         """Get the owner address of the given entity."""
         owner_metadata = metadata.get("owner")