Add a context manager for primitive callbacks

runtimepy/primitives/base.py

@@ -3,11 +3,13 @@
 """
 
 # built-in
+from contextlib import contextmanager as _contextmanager
 from copy import copy as _copy
 from math import isclose as _isclose
 from typing import BinaryIO as _BinaryIO
 from typing import Callable as _Callable
 from typing import Generic as _Generic
+from typing import Iterator as _Iterator
 from typing import TypeVar as _TypeVar
 
 # third-party
@@ -104,6 +106,18 @@ def remove_callback(self, callback_id: int) -> bool:
             del self.callbacks[callback_id]
         return result
 
+    @_contextmanager
+    def callback(
+        self, callback: PrimitiveChangeCallaback[T]
+    ) -> _Iterator[None]:
+        """Register a callback as a managed context."""
+
+        ident = self.register_callback(callback)
+        try:
+            yield
+        finally:
+            self.remove_callback(ident)
+
     @property
     def value(self) -> T:
         """Obtain the underlying value."""

tests/primitives/test_bool.py

@@ -17,3 +17,19 @@ def test_bool_basic():
 
     prim.clear()
     assert copied()
+
+    call_count = 0
+
+    def change_cb(_: bool, __: bool) -> None:
+        """A sample callback."""
+        nonlocal call_count
+        call_count += 1
+
+    with prim.callback(change_cb):
+        prim.toggle()
+        prim.toggle()
+
+    prim.toggle()
+    prim.toggle()
+
+    assert call_count == 2