Add owned registry subscriptions

Author: trissim

src/objectstate/__init__.py

@@ -222,7 +222,10 @@
 
 
 # Object state management
-from objectstate.object_state_registry import ObjectStateRegistry
+from objectstate.object_state_registry import (
+    ObjectStateRegistry,
+    ObjectStateRegistrySubscription,
+)
 from objectstate.object_state import ObjectState
 from objectstate.object_state_metadata import (
     ObjectStateMetadataContract,
@@ -259,6 +262,7 @@
     # Object state
     'ObjectState',
     'ObjectStateRegistry',
+    'ObjectStateRegistrySubscription',
     'ObjectStateMetadataContract',
     'ObjectStateMetadataContractRegistry',
     'ObjectStateMetadataStore',

src/objectstate/object_state_registry.py

@@ -5,7 +5,19 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, is_dataclass
 import logging
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TYPE_CHECKING, Generator, TypeAlias
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    TypeAlias,
+    TypeVar,
+)
 import copy
 
 from objectstate.object_state_metadata import ObjectStateMetadataStore
@@ -18,9 +30,26 @@
 
 TimeTravelStateEntry: TypeAlias = Tuple[str, 'ObjectState']
 TimeTravelCompleteCallback: TypeAlias = Callable[[List[TimeTravelStateEntry], Optional[str]], None]
+RegistryCallback = TypeVar("RegistryCallback", bound=Callable[..., None])
 logger = logging.getLogger(__name__)
 
 
+class ObjectStateRegistrySubscription:
+    """Idempotent ownership handle for one registry callback registration."""
+
+    def __init__(self, release: Callable[[], bool] | None) -> None:
+        self._release = release
+
+    def release(self) -> bool:
+        """Release this registration once, returning whether it was active."""
+
+        release = self._release
+        if release is None:
+            return False
+        self._release = None
+        return release()
+
+
 @dataclass(frozen=True)
 class TimeTravelTransaction:
     """Nominal request to move the registry to one snapshot."""
@@ -139,57 +168,96 @@ class ObjectStateRegistry:
     _on_history_changed_callbacks: List[Callable[[], None]] = []
 
     @classmethod
-    def add_register_callback(cls, callback: Callable[[str, 'ObjectState'], None]) -> None:
-        """Subscribe to ObjectState registration events."""
-        if callback not in cls._on_register_callbacks:
-            cls._on_register_callbacks.append(callback)
+    def add_register_callback(
+        cls,
+        callback: Callable[[str, 'ObjectState'], None],
+    ) -> ObjectStateRegistrySubscription:
+        """Subscribe to registration events and return its ownership handle."""
+        return cls._subscribe_callback(cls._on_register_callbacks, callback)
 
     @classmethod
-    def remove_register_callback(cls, callback: Callable[[str, 'ObjectState'], None]) -> None:
+    def remove_register_callback(
+        cls,
+        callback: Callable[[str, 'ObjectState'], None],
+    ) -> bool:
         """Unsubscribe from ObjectState registration events."""
-        if callback in cls._on_register_callbacks:
-            cls._on_register_callbacks.remove(callback)
+        return cls._remove_callback(cls._on_register_callbacks, callback)
 
     @classmethod
-    def add_unregister_callback(cls, callback: Callable[[str, 'ObjectState'], None]) -> None:
-        """Subscribe to ObjectState unregistration events."""
-        if callback not in cls._on_unregister_callbacks:
-            cls._on_unregister_callbacks.append(callback)
+    def add_unregister_callback(
+        cls,
+        callback: Callable[[str, 'ObjectState'], None],
+    ) -> ObjectStateRegistrySubscription:
+        """Subscribe to unregistration events and return its ownership handle."""
+        return cls._subscribe_callback(cls._on_unregister_callbacks, callback)
 
     @classmethod
-    def remove_unregister_callback(cls, callback: Callable[[str, 'ObjectState'], None]) -> None:
+    def remove_unregister_callback(
+        cls,
+        callback: Callable[[str, 'ObjectState'], None],
+    ) -> bool:
         """Unsubscribe from ObjectState unregistration events."""
-        if callback in cls._on_unregister_callbacks:
-            cls._on_unregister_callbacks.remove(callback)
+        return cls._remove_callback(cls._on_unregister_callbacks, callback)
 
     @classmethod
-    def add_time_travel_complete_callback(cls, callback: TimeTravelCompleteCallback) -> None:
+    def add_time_travel_complete_callback(
+        cls,
+        callback: TimeTravelCompleteCallback,
+    ) -> ObjectStateRegistrySubscription:
         """Subscribe to time-travel completion events.
 
         Callback receives (dirty_states, triggering_scope) where:
         - dirty_states: list of (scope_id, ObjectState) tuples with unsaved changes
         - triggering_scope: scope_id that triggered the snapshot (may be None)
         """
-        if callback not in cls._on_time_travel_complete_callbacks:
-            cls._on_time_travel_complete_callbacks.append(callback)
+        return cls._subscribe_callback(
+            cls._on_time_travel_complete_callbacks,
+            callback,
+        )
 
     @classmethod
-    def remove_time_travel_complete_callback(cls, callback: TimeTravelCompleteCallback) -> None:
+    def remove_time_travel_complete_callback(
+        cls,
+        callback: TimeTravelCompleteCallback,
+    ) -> bool:
         """Unsubscribe from time-travel completion events."""
-        if callback in cls._on_time_travel_complete_callbacks:
-            cls._on_time_travel_complete_callbacks.remove(callback)
+        return cls._remove_callback(cls._on_time_travel_complete_callbacks, callback)
 
     @classmethod
-    def add_history_changed_callback(cls, callback: Callable[[], None]) -> None:
-        """Subscribe to history change events (snapshot added or time-travel)."""
-        if callback not in cls._on_history_changed_callbacks:
-            cls._on_history_changed_callbacks.append(callback)
+    def add_history_changed_callback(
+        cls,
+        callback: Callable[[], None],
+    ) -> ObjectStateRegistrySubscription:
+        """Subscribe to history changes and return its ownership handle."""
+        return cls._subscribe_callback(cls._on_history_changed_callbacks, callback)
 
     @classmethod
-    def remove_history_changed_callback(cls, callback: Callable[[], None]) -> None:
+    def remove_history_changed_callback(cls, callback: Callable[[], None]) -> bool:
         """Unsubscribe from history change events."""
-        if callback in cls._on_history_changed_callbacks:
-            cls._on_history_changed_callbacks.remove(callback)
+        return cls._remove_callback(cls._on_history_changed_callbacks, callback)
+
+    @classmethod
+    def _subscribe_callback(
+        cls,
+        callbacks: List[RegistryCallback],
+        callback: RegistryCallback,
+    ) -> ObjectStateRegistrySubscription:
+        if callback in callbacks:
+            return ObjectStateRegistrySubscription(None)
+        callbacks.append(callback)
+        return ObjectStateRegistrySubscription(
+            lambda: cls._remove_callback(callbacks, callback)
+        )
+
+    @staticmethod
+    def _remove_callback(
+        callbacks: List[RegistryCallback],
+        callback: RegistryCallback,
+    ) -> bool:
+        if callback not in callbacks:
+            return False
+        callbacks.remove(callback)
+        return True
 
     @classmethod
     def _fire_history_changed_callbacks(cls) -> None:
@@ -474,39 +542,42 @@ def _notify_change(cls) -> None:
             cls._change_callbacks.remove(cb)
 
     @classmethod
-    def connect_listener(cls, callback: Callable[[], None]) -> None:
+    def connect_listener(
+        cls,
+        callback: Callable[[], None],
+    ) -> ObjectStateRegistrySubscription:
         """Connect a listener callback that's called on any change.
 
         The callback should debounce and call collect() to get fresh values.
         """
-        if callback not in cls._change_callbacks:
-            cls._change_callbacks.append(callback)
+        subscription = cls._subscribe_callback(cls._change_callbacks, callback)
+        if callback in cls._change_callbacks:
             logger.debug(f"Connected change listener: {callback}")
+        return subscription
 
     @classmethod
-    def disconnect_listener(cls, callback: Callable[[], None]) -> None:
+    def disconnect_listener(cls, callback: Callable[[], None]) -> bool:
         """Disconnect a change listener."""
-        if callback in cls._change_callbacks:
-            cls._change_callbacks.remove(callback)
+        removed = cls._remove_callback(cls._change_callbacks, callback)
+        if removed:
             logger.debug(f"Disconnected change listener: {callback}")
+        return removed
 
     @classmethod
     def add_resolved_changed_callback(
         cls,
         callback: Callable[[str, Set[str]], None],
-    ) -> None:
-        """Subscribe to resolved-value changes from any registered ObjectState."""
-        if callback not in cls._resolved_changed_callbacks:
-            cls._resolved_changed_callbacks.append(callback)
+    ) -> ObjectStateRegistrySubscription:
+        """Subscribe to resolved changes and return its ownership handle."""
+        return cls._subscribe_callback(cls._resolved_changed_callbacks, callback)
 
     @classmethod
     def remove_resolved_changed_callback(
         cls,
         callback: Callable[[str, Set[str]], None],
-    ) -> None:
+    ) -> bool:
         """Unsubscribe from registry-wide resolved-value changes."""
-        if callback in cls._resolved_changed_callbacks:
-            cls._resolved_changed_callbacks.remove(callback)
+        return cls._remove_callback(cls._resolved_changed_callbacks, callback)
 
     @classmethod
     def notify_resolved_changed(

tests/test_registry_subscription.py

@@ -0,0 +1,88 @@
+"""Ownership contracts for ObjectState registry callback subscriptions."""
+
+import pytest
+
+from objectstate import ObjectStateRegistry, ObjectStateRegistrySubscription
+
+
+@pytest.mark.parametrize(
+    ("add", "remove", "callback"),
+    (
+        (
+            ObjectStateRegistry.add_register_callback,
+            ObjectStateRegistry.remove_register_callback,
+            lambda _scope, _state: None,
+        ),
+        (
+            ObjectStateRegistry.add_unregister_callback,
+            ObjectStateRegistry.remove_unregister_callback,
+            lambda _scope, _state: None,
+        ),
+        (
+            ObjectStateRegistry.add_time_travel_complete_callback,
+            ObjectStateRegistry.remove_time_travel_complete_callback,
+            lambda _states, _scope: None,
+        ),
+        (
+            ObjectStateRegistry.add_history_changed_callback,
+            ObjectStateRegistry.remove_history_changed_callback,
+            lambda: None,
+        ),
+        (
+            ObjectStateRegistry.add_resolved_changed_callback,
+            ObjectStateRegistry.remove_resolved_changed_callback,
+            lambda _scope, _paths: None,
+        ),
+        (
+            ObjectStateRegistry.connect_listener,
+            ObjectStateRegistry.disconnect_listener,
+            lambda: None,
+        ),
+    ),
+)
+def test_registry_callback_axes_return_one_release_handle(
+    add,
+    remove,
+    callback,
+) -> None:
+    subscription = add(callback)
+    try:
+        assert isinstance(subscription, ObjectStateRegistrySubscription)
+        assert subscription.release()
+        assert not subscription.release()
+    finally:
+        remove(callback)
+
+
+def test_history_callback_subscription_releases_exact_registration() -> None:
+    observed = []
+
+    def observe() -> None:
+        observed.append(True)
+
+    subscription = ObjectStateRegistry.add_history_changed_callback(observe)
+    try:
+        assert isinstance(subscription, ObjectStateRegistrySubscription)
+        ObjectStateRegistry._fire_history_changed_callbacks()
+        assert observed == [True]
+
+        assert subscription.release()
+        assert not subscription.release()
+        ObjectStateRegistry._fire_history_changed_callbacks()
+        assert observed == [True]
+    finally:
+        ObjectStateRegistry.remove_history_changed_callback(observe)
+
+
+def test_duplicate_registration_does_not_transfer_existing_ownership() -> None:
+    def observe() -> None:
+        pass
+
+    owner = ObjectStateRegistry.add_history_changed_callback(observe)
+    duplicate = ObjectStateRegistry.add_history_changed_callback(observe)
+    try:
+        assert not duplicate.release()
+        ObjectStateRegistry._fire_history_changed_callbacks()
+        assert owner.release()
+    finally:
+        ObjectStateRegistry.remove_history_changed_callback(observe)