Add AttrR wait_for_value

Author: GDYendell

docs/conf.py

@@ -77,6 +77,7 @@
     ("py:class", "asyncio.events.AbstractEventLoop"),
     ("py:class", "asyncio.streams.StreamReader"),
     ("py:class", "asyncio.streams.StreamWriter"),
+    ("py:class", "asyncio.locks.Event"),
     # Annoying error:
     #     docstring of collections.abc.Callable:1: WARNING:
     #     'any' reference target not found: self [ref.any]

src/fastcs/attributes/attr_r.py

@@ -6,6 +6,7 @@
 
 from fastcs.attributes.attribute import Attribute
 from fastcs.attributes.attribute_io_ref import AttributeIORefT
+from fastcs.attributes.util import AttrValuePredicate, PredicateEvent
 from fastcs.datatypes import DataType, DType_T
 from fastcs.logging import bind_logger
 
@@ -39,6 +40,8 @@ def __init__(
         """Callback to update the value of the attribute with an IO to the source"""
         self._on_update_callbacks: list[AttrOnUpdateCallback[DType_T]] | None = None
         """Callbacks to publish changes to the value of the attribute"""
+        self._on_update_events: set[PredicateEvent[DType_T]] = set()
+        """Events to set when the value satisifies some predicate"""
 
     def get(self) -> DType_T:
         """Get the cached value of the attribute."""
@@ -67,6 +70,10 @@ async def update(self, value: Any) -> None:
 
         self._value = self._datatype.validate(value)
 
+        self._on_update_events -= {
+            e for e in self._on_update_events if e.set(self._value)
+        }
+
         if self._on_update_callbacks is not None:
             try:
                 await asyncio.gather(
@@ -78,6 +85,16 @@ async def update(self, value: Any) -> None:
                 )
                 raise
 
+    def _register_update_event(self, event: PredicateEvent):
+        """Register an event to be set when the value satisfies a predicate
+
+        Args:
+            predicate: The predicate to filter event set by
+            event: The event to set
+
+        """
+        self._on_update_events.add(event)
+
     def add_on_update_callback(self, callback: AttrOnUpdateCallback[DType_T]) -> None:
         """Add a callback to be called when the value of the attribute is updated
 
@@ -115,3 +132,62 @@ async def update_attribute():
                 raise
 
         return update_attribute
+
+    async def wait_for_predicate(
+        self, predicate: AttrValuePredicate[DType_T], *, timeout: float
+    ):
+        """Wait for the predicate to be satisfied when called with the current value
+
+        Args:
+            predicate: The predicate to test - a callable that takes the attribute
+                value and returns True if the event should be set
+            timeout: The timeout in seconds
+
+        """
+        if predicate(self._value):
+            self.log_event(
+                "Predicate already satisfied", predicate=predicate, attribute=self
+            )
+            return
+
+        self._register_update_event(update_event := PredicateEvent(predicate))
+
+        self.log_event("Waiting for predicate", predicate=predicate, attribute=self)
+        try:
+            await asyncio.wait_for(update_event.wait(), timeout)
+        except TimeoutError:
+            self._on_update_events.remove(update_event)
+            raise TimeoutError(
+                f"Timeout waiting for predicate {predicate}. "
+                f"Current value: {self._value}"
+            ) from None
+
+        self.log_event("Predicate satisfied", predicate=predicate, attribute=self)
+
+    async def wait_for_value(self, value: DType_T, *, timeout: float):
+        """Wait for value to equal the required value
+
+        Args:
+            value: The value to wait for
+            timeout: The timeout in seconds
+
+        Raises:
+            TimeoutError: If the attribute does not reach the required value within the
+                timeout
+
+        """
+        if self._value == value:
+            self.log_event("Value already equal", value=value, attribute=self)
+            return
+
+        def predicate(v: DType_T) -> bool:
+            return v == value
+
+        try:
+            await self.wait_for_predicate(predicate, timeout=timeout)
+        except TimeoutError:
+            raise TimeoutError(
+                f"Timeout waiting for value {value}. Current value: {self._value}"
+            ) from None
+
+        self.log_event("Value equal", value=value, attribute=self)

src/fastcs/attributes/util.py

@@ -0,0 +1,39 @@
+import asyncio
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Generic
+
+from fastcs.datatypes import DType_T
+
+AttrValuePredicate = Callable[[DType_T], bool]
+
+
+@dataclass(eq=False)
+class PredicateEvent(Generic[DType_T]):
+    """A wrapper of an asyncio.Event that only triggers when a predicate is satisfied"""
+
+    _predicate: AttrValuePredicate[DType_T]
+    """Predicate to filter set calls by"""
+    _event: asyncio.Event = field(default_factory=asyncio.Event)
+    """Event to set"""
+
+    def set(self, value: DType_T):
+        """Set the event if the predicate is satisfied by the value
+
+        Returns:
+            `True` if the event was set, else `False`
+
+        """
+        if self._predicate(value):
+            self._event.set()
+            return True
+
+        return False
+
+    async def wait(self):
+        """Wait for the event to be set"""
+        await self._event.wait()
+
+    def __hash__(self) -> int:
+        """Make instances unique when stored in sets"""
+        return id(self)

tests/test_attributes.py

@@ -1,3 +1,4 @@
+import asyncio
 from dataclasses import dataclass
 from functools import partial
 from typing import Generic, TypeVar
@@ -12,7 +13,7 @@
 NumberT = TypeVar("NumberT", int, float)
 
 
-def test_attribute():
+def test_attr_r():
     attr = AttrR(String(), group="test group")
 
     with pytest.raises(RuntimeError):
@@ -39,6 +40,57 @@ def test_attribute():
     assert attr.get() == ""
 
 
+@pytest.mark.asyncio
+async def test_wait_for_predicate(mocker: MockerFixture):
+    attr = AttrR(Int(), initial_value=0)
+
+    async def update(attr: AttrR):
+        while True:
+            await asyncio.sleep(0.1)
+            await attr.update(attr.get() + 3)  # 3, 6, 9, 12 != 10
+
+    asyncio.create_task(update(attr))
+
+    # We won't see exactly 10 so check for greater than
+    def predicate(v: int) -> bool:
+        return v > 10
+
+    register_spy = mocker.spy(attr, "_register_update_event")
+    with pytest.raises(TimeoutError):
+        await attr.wait_for_predicate(predicate, timeout=0.2)
+
+    await attr.wait_for_predicate(predicate, timeout=1)
+
+    assert register_spy.call_count == 2
+
+    # Returns immediately without creating event if value already as expected
+    await attr.wait_for_predicate(predicate, timeout=1)
+    assert register_spy.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_wait_for_value(mocker: MockerFixture):
+    attr = AttrR(Int(), initial_value=0)
+
+    async def update(attr: AttrR):
+        await asyncio.sleep(0.5)
+        await attr.update(1)
+
+    asyncio.create_task(update(attr))
+
+    register_spy = mocker.spy(attr, "_register_update_event")
+    with pytest.raises(TimeoutError):
+        await attr.wait_for_value(1, timeout=0.2)
+
+    await attr.wait_for_value(1, timeout=1)
+
+    assert register_spy.call_count == 2
+
+    # Returns immediately without creating event if value already as expected
+    await attr.wait_for_value(1, timeout=1)
+    assert register_spy.call_count == 2
+
+
 @pytest.mark.asyncio
 async def test_attributes():
     device = {"state": "Idle", "number": 1, "count": False}