ophyd-async 0.8.0a5__py3-none-any.whl → 0.9.0__py3-none-any.whl

ophyd_async/core/_signal.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import asyncio
 import functools
-from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping
+import time
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import Any, Generic, cast
 
 from bluesky.protocols import (
@@ -17,7 +18,6 @@ from event_model import DataKey
 from ._device import Device, DeviceConnector
 from ._mock_signal_backend import MockSignalBackend
 from ._protocol import (
-    AsyncConfigurable,
     AsyncReadable,
     AsyncStageable,
     Reading,
@@ -28,7 +28,7 @@ from ._signal_backend import (
     SignalDatatypeV,
 )
 from ._soft_signal_backend import SoftSignalBackend
-from ._status import AsyncStatus
+from ._status import AsyncStatus, completed_status
 from ._utils import (
     CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
@@ -97,32 +97,37 @@ class Signal(Device, Generic[SignalDatatypeT]):
 
 
 class _SignalCache(Generic[SignalDatatypeT]):
-    def __init__(self, backend: SignalBackend[SignalDatatypeT], signal: Signal):
-        self._signal = signal
+    def __init__(self, backend: SignalBackend[SignalDatatypeT], signal: Signal) -> None:
+        self._signal: Signal[Any] = signal
         self._staged = False
         self._listeners: dict[Callback, bool] = {}
         self._valid = asyncio.Event()
         self._reading: Reading[SignalDatatypeT] | None = None
-        self.backend = backend
+        self.backend: SignalBackend[SignalDatatypeT] = backend
         signal.log.debug(f"Making subscription on source {signal.source}")
         backend.set_callback(self._callback)
 
-    def close(self):
+    def close(self) -> None:
         self.backend.set_callback(None)
         self._signal.log.debug(f"Closing subscription on source {self._signal.source}")
 
+    def _ensure_reading(self) -> Reading[SignalDatatypeT]:
+        if not self._reading:
+            msg = "Monitor not working"
+            raise RuntimeError(msg)
+        return self._reading
+
     async def get_reading(self) -> Reading[SignalDatatypeT]:
         await self._valid.wait()
-        assert self._reading is not None, "Monitor not working"
-        return self._reading
+        return self._ensure_reading()
 
     async def get_value(self) -> SignalDatatypeT:
-        reading = await self.get_reading()
+        reading: Reading[SignalDatatypeT] = await self.get_reading()
         return reading["value"]
 
-    def _callback(self, reading: Reading[SignalDatatypeT]):
+    def _callback(self, reading: Reading[SignalDatatypeT]) -> None:
         self._signal.log.debug(
-            f"Updated subscription: reading of source {self._signal.source} changed"
+            f"Updated subscription: reading of source {self._signal.source} changed "
             f"from {self._reading} to {reading}"
         )
         self._reading = reading
@@ -134,12 +139,10 @@ class _SignalCache(Generic[SignalDatatypeT]):
         self,
         function: Callback[dict[str, Reading[SignalDatatypeT]] | SignalDatatypeT],
         want_value: bool,
-    ):
-        assert self._reading, "Monitor not working"
-        if want_value:
-            function(self._reading["value"])
-        else:
-            function({self._signal.name: self._reading})
+    ) -> None:
+        function(self._ensure_reading()["value"]) if want_value else function(
+            {self._signal.name: self._ensure_reading()}
+        )
 
     def subscribe(self, function: Callback, want_value: bool) -> None:
         self._listeners[function] = want_value
@@ -150,7 +153,7 @@ class _SignalCache(Generic[SignalDatatypeT]):
         self._listeners.pop(function)
         return self._staged or bool(self._listeners)
 
-    def set_staged(self, staged: bool):
+    def set_staged(self, staged: bool) -> bool:
         self._staged = staged
         return self._staged or bool(self._listeners)
 
@@ -167,7 +170,10 @@ class SignalR(Signal[SignalDatatypeT], AsyncReadable, AsyncStageable, Subscribab
         if cached is None:
             cached = self._cache is not None
         if cached:
-            assert self._cache, f"{self.source} not being monitored"
+            if not self._cache:
+                msg = f"{self.source} not being monitored"
+                raise RuntimeError(msg)
+            # assert self._cache, f"{self.source} not being monitored"
             return self._cache
         else:
             return self._connector.backend
@@ -301,137 +307,70 @@ def soft_signal_r_and_setter(
     return (signal, backend.set_value)
 
 
-def _generate_assert_error_msg(name: str, expected_result, actual_result) -> str:
-    WARNING = "\033[93m"
-    FAIL = "\033[91m"
-    ENDC = "\033[0m"
-    return (
-        f"Expected {WARNING}{name}{ENDC} to produce"
-        + f"\n{FAIL}{expected_result}{ENDC}"
-        + f"\nbut actually got \n{FAIL}{actual_result}{ENDC}"
-    )
-
-
-async def assert_value(signal: SignalR[SignalDatatypeT], value: Any) -> None:
-    """Assert a signal's value and compare it an expected signal.
+async def observe_value(
+    signal: SignalR[SignalDatatypeT],
+    timeout: float | None = None,
+    done_status: Status | None = None,
+    done_timeout: float | None = None,
+) -> AsyncGenerator[SignalDatatypeT, None]:
+    """Subscribe to the value of a signal so it can be iterated from.
 
     Parameters
     ----------
     signal:
-        signal with get_value.
-    value:
-        The expected value from the signal.
+        Call subscribe_value on this at the start, and clear_sub on it at the
+        end
+    timeout:
+        If given, how long to wait for each updated value in seconds. If an update
+        is not produced in this time then raise asyncio.TimeoutError
+    done_status:
+        If this status is complete, stop observing and make the iterator return.
+        If it raises an exception then this exception will be raised by the iterator.
+    done_timeout:
+        If given, the maximum time to watch a signal, in seconds. If the loop is still
+        being watched after this length, raise asyncio.TimeoutError. This should be used
+        instead of on an 'asyncio.wait_for' timeout
 
     Notes
     -----
-    Example usage::
-        await assert_value(signal, value)
+    Due to a rare condition with busy signals, it is not recommended to use this
+    function with asyncio.timeout, including in an 'asyncio.wait_for' loop. Instead,
+    this timeout should be given to the done_timeout parameter.
 
-    """
-    actual_value = await signal.get_value()
-    assert actual_value == value, _generate_assert_error_msg(
-        name=signal.name,
-        expected_result=value,
-        actual_result=actual_value,
-    )
-
-
-async def assert_reading(
-    readable: AsyncReadable, expected_reading: Mapping[str, Reading]
-) -> None:
-    """Assert readings from readable.
-
-    Parameters
-    ----------
-    readable:
-        Callable with readable.read function that generate readings.
-
-    reading:
-        The expected readings from the readable.
-
-    Notes
-    -----
     Example usage::
-        await assert_reading(readable, reading)
-
-    """
-    actual_reading = await readable.read()
-    assert expected_reading == actual_reading, _generate_assert_error_msg(
-        name=readable.name,
-        expected_result=expected_reading,
-        actual_result=actual_reading,
-    )
-
-
-async def assert_configuration(
-    configurable: AsyncConfigurable,
-    configuration: Mapping[str, Reading],
-) -> None:
-    """Assert readings from Configurable.
-
-    Parameters
-    ----------
-    configurable:
-        Configurable with Configurable.read function that generate readings.
-
-    configuration:
-        The expected readings from configurable.
-
-    Notes
-    -----
-    Example usage::
-        await assert_configuration(configurable configuration)
 
+        async for value in observe_value(sig):
+            do_something_with(value)
     """
-    actual_configurable = await configurable.read_configuration()
-    assert configuration == actual_configurable, _generate_assert_error_msg(
-        name=configurable.name,
-        expected_result=configuration,
-        actual_result=actual_configurable,
-    )
-
 
-def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int):
-    """Assert emitted document generated by running a Bluesky plan
-
-    Parameters
-    ----------
-    Doc:
-        A dictionary
+    async for _, value in observe_signals_value(
+        signal,
+        timeout=timeout,
+        done_status=done_status,
+        done_timeout=done_timeout,
+    ):
+        yield value
 
-    numbers:
-        expected emission in kwarg from
 
-    Notes
-    -----
-    Example usage::
-        assert_emitted(docs, start=1, descriptor=1,
-        resource=1, datum=1, event=1, stop=1)
-    """
-    assert list(docs) == list(numbers), _generate_assert_error_msg(
-        name="documents",
-        expected_result=list(numbers),
-        actual_result=list(docs),
-    )
-    actual_numbers = {name: len(d) for name, d in docs.items()}
-    assert actual_numbers == numbers, _generate_assert_error_msg(
-        name="emitted",
-        expected_result=numbers,
-        actual_result=actual_numbers,
-    )
+def _get_iteration_timeout(
+    timeout: float | None, overall_deadline: float | None
+) -> float | None:
+    overall_deadline = overall_deadline - time.monotonic() if overall_deadline else None
+    return min([x for x in [overall_deadline, timeout] if x is not None], default=None)
 
 
-async def observe_value(
-    signal: SignalR[SignalDatatypeT],
+async def observe_signals_value(
+    *signals: SignalR[SignalDatatypeT],
     timeout: float | None = None,
     done_status: Status | None = None,
-) -> AsyncGenerator[SignalDatatypeT, None]:
+    done_timeout: float | None = None,
+) -> AsyncGenerator[tuple[SignalR[SignalDatatypeT], SignalDatatypeT], None]:
     """Subscribe to the value of a signal so it can be iterated from.
 
     Parameters
     ----------
-    signal:
-        Call subscribe_value on this at the start, and clear_sub on it at the
+    signals:
+        Call subscribe_value on all the signals at the start, and clear_sub on it at the
         end
     timeout:
         If given, how long to wait for each updated value in seconds. If an update
@@ -439,36 +378,57 @@ async def observe_value(
     done_status:
         If this status is complete, stop observing and make the iterator return.
         If it raises an exception then this exception will be raised by the iterator.
+    done_timeout:
+        If given, the maximum time to watch a signal, in seconds. If the loop is still
+        being watched after this length, raise asyncio.TimeoutError. This should be used
+        instead of on an 'asyncio.wait_for' timeout
 
     Notes
     -----
     Example usage::
 
-        async for value in observe_value(sig):
-            do_something_with(value)
+        async for signal,value in observe_signals_values(sig1,sig2,..):
+            if signal is sig1:
+                do_something_with(value)
+            elif signal is sig2:
+                do_something_else_with(value)
     """
+    q: asyncio.Queue[tuple[SignalR[SignalDatatypeT], SignalDatatypeT] | Status] = (
+        asyncio.Queue()
+    )
+
+    cbs: dict[SignalR, Callback] = {}
+    for signal in signals:
 
-    q: asyncio.Queue[SignalDatatypeT | Status] = asyncio.Queue()
+        def queue_value(value: SignalDatatypeT, signal=signal):
+            q.put_nowait((signal, value))
+
+        cbs[signal] = queue_value
+        signal.subscribe_value(queue_value)
 
     if done_status is not None:
         done_status.add_callback(q.put_nowait)
-
-    signal.subscribe_value(q.put_nowait)
+    overall_deadline = time.monotonic() + done_timeout if done_timeout else None
     try:
         while True:
-            # yield here in case something else is filling the queue
-            # like in test_observe_value_times_out_with_no_external_task()
-            await asyncio.sleep(0)
-            item = await asyncio.wait_for(q.get(), timeout)
+            if overall_deadline and time.monotonic() >= overall_deadline:
+                raise asyncio.TimeoutError(
+                    f"observe_value was still observing signals "
+                    f"{[signal.source for signal in signals]} after "
+                    f"timeout {done_timeout}s"
+                )
+            iteration_timeout = _get_iteration_timeout(timeout, overall_deadline)
+            item = await asyncio.wait_for(q.get(), iteration_timeout)
             if done_status and item is done_status:
                 if exc := done_status.exception():
                     raise exc
                 else:
                     break
             else:
-                yield cast(SignalDatatypeT, item)
+                yield cast(tuple[SignalR[SignalDatatypeT], SignalDatatypeT], item)
     finally:
-        signal.clear_sub(q.put_nowait)
+        for signal, cb in cbs.items():
+            signal.clear_sub(cb)
 
 
 class _ValueChecker(Generic[SignalDatatypeT]):
@@ -533,15 +493,16 @@ async def wait_for_value(
 async def set_and_wait_for_other_value(
     set_signal: SignalW[SignalDatatypeT],
     set_value: SignalDatatypeT,
-    read_signal: SignalR[SignalDatatypeV],
-    read_value: SignalDatatypeV,
+    match_signal: SignalR[SignalDatatypeV],
+    match_value: SignalDatatypeV | Callable[[SignalDatatypeV], bool],
     timeout: float = DEFAULT_TIMEOUT,
     set_timeout: float | None = None,
+    wait_for_set_completion: bool = True,
 ) -> AsyncStatus:
     """Set a signal and monitor another signal until it has the specified value.
 
     This function sets a set_signal to a specified set_value and waits for
-    a read_signal to have the read_value.
+    a match_signal to have the match_value.
 
     Parameters
     ----------
@@ -549,14 +510,16 @@ async def set_and_wait_for_other_value(
         The signal to set
     set_value:
         The value to set it to
-    read_signal:
+    match_signal:
         The signal to monitor
-    read_value:
+    match_value:
         The value to wait for
     timeout:
         How long to wait for the signal to have the value
     set_timeout:
         How long to wait for the set to complete
+    wait_for_set_completion:
+        This will wait for set completion #More info in how-to docs
 
     Notes
     -----
@@ -565,7 +528,7 @@ async def set_and_wait_for_other_value(
         set_and_wait_for_value(device.acquire, 1, device.acquire_rbv, 1)
     """
     # Start monitoring before the set to avoid a race condition
-    values_gen = observe_value(read_signal)
+    values_gen = observe_value(match_signal)
 
     # Get the initial value from the monitor to make sure we've created it
     current_value = await anext(values_gen)
@@ -573,28 +536,33 @@ async def set_and_wait_for_other_value(
     status = set_signal.set(set_value, timeout=set_timeout)
 
     # If the value was the same as before no need to wait for it to change
-    if current_value != read_value:
+    if current_value != match_value:
 
         async def _wait_for_value():
             async for value in values_gen:
-                if value == read_value:
+                if value == match_value:
                     break
 
         try:
             await asyncio.wait_for(_wait_for_value(), timeout)
+            if wait_for_set_completion:
+                await status
+            return status
         except asyncio.TimeoutError as e:
             raise TimeoutError(
-                f"{read_signal.name} didn't match {read_value} in {timeout}s"
+                f"{match_signal.name} didn't match {match_value} in {timeout}s"
             ) from e
 
-    return status
+    return completed_status()
 
 
 async def set_and_wait_for_value(
     signal: SignalRW[SignalDatatypeT],
     value: SignalDatatypeT,
+    match_value: SignalDatatypeT | Callable[[SignalDatatypeT], bool] | None = None,
     timeout: float = DEFAULT_TIMEOUT,
     status_timeout: float | None = None,
+    wait_for_set_completion: bool = True,
 ) -> AsyncStatus:
     """Set a signal and monitor it until it has that value.
 
@@ -609,10 +577,15 @@ async def set_and_wait_for_value(
         The signal to set
     value:
         The value to set it to
+    match_value:
+        The expected value of the signal after the operation.
+        Used to verify that the set operation was successful.
     timeout:
         How long to wait for the signal to have the value
     status_timeout:
         How long the returned Status will wait for the set to complete
+    wait_for_set_completion:
+        This will wait for set completion #More info in how-to docs
 
     Notes
     -----
@@ -620,6 +593,46 @@ async def set_and_wait_for_value(
 
         set_and_wait_for_value(device.acquire, 1)
     """
+    if match_value is None:
+        match_value = value
     return await set_and_wait_for_other_value(
-        signal, value, signal, value, timeout, status_timeout
+        signal,
+        value,
+        signal,
+        match_value,
+        timeout,
+        status_timeout,
+        wait_for_set_completion,
     )
+
+
+def walk_rw_signals(device: Device, path_prefix: str = "") -> dict[str, SignalRW[Any]]:
+    """Retrieve all SignalRWs from a device.
+
+    Stores retrieved signals with their dotted attribute paths in a dictionary. Used as
+    part of saving and loading a device.
+
+    Parameters
+    ----------
+    device : Device
+        Ophyd device to retrieve read-write signals from.
+
+    path_prefix : str
+        For internal use, leave blank when calling the method.
+
+    Returns
+    -------
+    SignalRWs : dict
+        A dictionary matching the string attribute path of a SignalRW with the
+        signal itself.
+
+    """
+    signals: dict[str, SignalRW[Any]] = {}
+
+    for attr_name, attr in device.children():
+        dot_path = f"{path_prefix}{attr_name}"
+        if type(attr) is SignalRW:
+            signals[dot_path] = attr
+        attr_signals = walk_rw_signals(attr, path_prefix=dot_path + ".")
+        signals.update(attr_signals)
+    return signals

ophyd_async/core/_signal_backend.py

@@ -10,7 +10,10 @@ from ._table import Table
 from ._utils import Callback, StrictEnum, T
 
 DTypeScalar_co = TypeVar("DTypeScalar_co", covariant=True, bound=np.generic)
-Array1D = np.ndarray[tuple[int], np.dtype[DTypeScalar_co]]
+# To be a 1D array shape should really be tuple[int], but np.array()
+# currently produces tuple[int, ...] even when it has 1D input args
+# https://github.com/numpy/numpy/issues/28077#issuecomment-2566485178
+Array1D = np.ndarray[tuple[int, ...], np.dtype[DTypeScalar_co]]
 Primitive = bool | int | float | str
 # NOTE: if you change this union then update the docs to match
 SignalDatatype = (

ophyd_async/core/_soft_signal_backend.py

@@ -175,7 +175,8 @@ class SoftSignalBackend(SignalBackend[SignalDatatypeT]):
         return self.reading["value"]
 
     def set_callback(self, callback: Callback[Reading[SignalDatatypeT]] | None) -> None:
+        if callback and self.callback:
+            raise RuntimeError("Cannot set a callback when one is already set")
         if callback:
-            assert not self.callback, "Cannot set a callback when one is already set"
             callback(self.reading)
         self.callback = callback

ophyd_async/core/_table.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from typing import Annotated, Any, TypeVar, get_origin
 
 import numpy as np
@@ -19,6 +19,13 @@ def _concat(value1, value2):
         return value1 + value2
 
 
+def _make_default_factory(dtype: np.dtype) -> Callable[[], np.ndarray]:
+    def numpy_array_default_factory() -> np.ndarray:
+        return np.array([], dtype)
+
+    return numpy_array_default_factory
+
+
 class Table(BaseModel):
     """An abstraction of a Table of str to numpy array."""
 
@@ -32,6 +39,11 @@ class Table(BaseModel):
     # so it is strictly checked against the BaseModel we are supplied.
     model_config = ConfigDict(extra="allow")
 
+    # Add an init method to match the above model config, otherwise the type
+    # checker will not think we can pass arbitrary kwargs into the base class init
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
     @classmethod
     def __init_subclass__(cls):
         # But forbit extra in subclasses so it gets validated
@@ -45,9 +57,7 @@ class Table(BaseModel):
                     NpArrayPydanticAnnotation.factory(
                         data_type=dtype.type, dimensions=1, strict_data_typing=False
                     ),
-                    Field(
-                        default_factory=lambda dtype=dtype: np.array([], dtype=dtype)
-                    ),
+                    Field(default_factory=_make_default_factory(dtype)),
                 ]
             elif get_origin(anno) is Sequence:
                 new_anno = Annotated[anno, Field(default_factory=list)]
@@ -73,9 +83,6 @@ class Table(BaseModel):
             }
         )
 
-    def __eq__(self, value: object) -> bool:
-        return super().__eq__(value)
-
     def numpy_dtype(self) -> np.dtype:
         dtype = []
         for k, v in self:
@@ -94,8 +101,10 @@ class Table(BaseModel):
                 v = v[selection]
             if array is None:
                 array = np.empty(v.shape, dtype=self.numpy_dtype())
-            array[k] = v
-        assert array is not None
+            array[k] = v  # type: ignore
+        if array is None:
+            msg = "No arrays found in table"
+            raise ValueError(msg)
         return array
 
     @model_validator(mode="before")
@@ -118,10 +127,12 @@ class Table(BaseModel):
                 # Convert to correct dtype, but only if we don't lose precision
                 # as a result
                 cast_value = np.array(data_value).astype(expected_dtype)
-                assert np.array_equal(data_value, cast_value), (
-                    f"{field_name}: Cannot cast {data_value} to {expected_dtype} "
-                    "without losing precision"
-                )
+                if not np.array_equal(data_value, cast_value):
+                    msg = (
+                        f"{field_name}: Cannot cast {data_value} to {expected_dtype} "
+                        "without losing precision"
+                    )
+                    raise ValueError(msg)
                 data_dict[field_name] = cast_value
         return data_dict
 
@@ -130,7 +141,9 @@ class Table(BaseModel):
         lengths: dict[int, set[str]] = {}
         for field_name, field_value in self:
             lengths.setdefault(len(field_value), set()).add(field_name)
-        assert len(lengths) <= 1, f"Columns should be same length, got {lengths=}"
+        if len(lengths) > 1:
+            msg = f"Columns should be same length, got {lengths=}"
+            raise ValueError(msg)
         return self
 
     def __len__(self) -> int: