ophyd-async 0.3a3__py3-none-any.whl → 0.3a5__py3-none-any.whl

ophyd_async/core/mock_signal_utils.py

@@ -0,0 +1,149 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, Callable, Generator, Iterable, Iterator, List
+from unittest.mock import ANY, Mock
+
+from ophyd_async.core.signal import Signal
+from ophyd_async.core.utils import T
+
+from .mock_signal_backend import MockSignalBackend
+
+
+def _get_mock_signal_backend(signal: Signal) -> MockSignalBackend:
+    assert isinstance(signal._backend, MockSignalBackend), (
+        "Expected to receive a `MockSignalBackend`, instead "
+        f" received {type(signal._backend)}. "
+    )
+    return signal._backend
+
+
+def set_mock_value(signal: Signal[T], value: T):
+    """Set the value of a signal that is in mock mode."""
+    backend = _get_mock_signal_backend(signal)
+    backend.set_value(value)
+
+
+def set_mock_put_proceeds(signal: Signal[T], proceeds: bool):
+    """Allow or block a put with wait=True from proceeding"""
+    backend = _get_mock_signal_backend(signal)
+
+    if proceeds:
+        backend.put_proceeds.set()
+    else:
+        backend.put_proceeds.clear()
+
+
+@asynccontextmanager
+async def mock_puts_blocked(*signals: List[Signal]):
+    for signal in signals:
+        set_mock_put_proceeds(signal, False)
+    yield
+    for signal in signals:
+        set_mock_put_proceeds(signal, True)
+
+
+def assert_mock_put_called_with(signal: Signal, value: Any, wait=ANY, timeout=ANY):
+    backend = _get_mock_signal_backend(signal)
+    backend.put_mock.assert_called_with(value, wait=wait, timeout=timeout)
+
+
+def reset_mock_put_calls(signal: Signal):
+    backend = _get_mock_signal_backend(signal)
+    backend.put_mock.reset_mock()
+
+
+class _SetValuesIterator:
+    # Garbage collected by the time __del__ is called unless we put it as a
+    # global attrbute here.
+    require_all_consumed: bool = False
+
+    def __init__(
+        self,
+        signal: Signal,
+        values: Iterable[Any],
+        require_all_consumed: bool = False,
+    ):
+        self.signal = signal
+        self.values = values
+        self.require_all_consumed = require_all_consumed
+        self.index = 0
+
+        self.iterator = enumerate(values, start=1)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        # Will propogate StopIteration
+        self.index, next_value = next(self.iterator)
+        set_mock_value(self.signal, next_value)
+        return next_value
+
+    def __del__(self):
+        if self.require_all_consumed and self.index != len(self.values):
+            raise AssertionError("Not all values have been consumed.")
+
+
+def set_mock_values(
+    signal: Signal,
+    values: Iterable[Any],
+    require_all_consumed: bool = False,
+) -> Iterator[Any]:
+    """Iterator to set a signal to a sequence of values, optionally repeating the
+    sequence.
+
+    Parameters
+    ----------
+    signal:
+        A signal with a `MockSignalBackend` backend.
+    values:
+        An iterable of the values to set the signal to, on each iteration
+        the value will be set.
+    require_all_consumed:
+        If True, an AssertionError will be raised if the iterator is deleted before
+        all values have been consumed.
+
+    Notes
+    -----
+    Example usage::
+
+           for value_set in set_mock_values(signal, [1, 2, 3]):
+                # do something
+
+            cm = set_mock_values(signal, 1, 2, 3, require_all_consumed=True):
+            next(cm)
+            # do something
+    """
+
+    return _SetValuesIterator(
+        signal,
+        values,
+        require_all_consumed=require_all_consumed,
+    )
+
+
+@contextmanager
+def _unset_side_effect_cm(put_mock: Mock):
+    yield
+    put_mock.side_effect = None
+
+
+# linting isn't smart enought to realize @contextmanager will give use a
+# ContextManager[None]
+def callback_on_mock_put(
+    signal: Signal, callback: Callable[[T], None]
+) -> Generator[None, None, None]:
+    """For setting a callback when a backend is put to.
+
+    Can either be used in a context, with the callback being
+    unset on exit, or as an ordinary function.
+
+    Parameters
+    ----------
+    signal:
+        A signal with a `MockSignalBackend` backend.
+    callback:
+        The callback to call when the backend is put to during the context.
+    """
+    backend = _get_mock_signal_backend(signal)
+    backend.put_mock.side_effect = callback
+    return _unset_side_effect_cm(backend.put_mock)

ophyd_async/core/signal.py

@@ -21,18 +21,18 @@ from bluesky.protocols import (
     Location,
     Movable,
     Reading,
+    Status,
     Subscribable,
 )
 
+from ophyd_async.core.mock_signal_backend import MockSignalBackend
 from ophyd_async.protocols import AsyncConfigurable, AsyncReadable, AsyncStageable
 
 from .async_status import AsyncStatus
 from .device import Device
 from .signal_backend import SignalBackend
-from .sim_signal_backend import SimSignalBackend
-from .utils import DEFAULT_TIMEOUT, Callback, ReadingValueCallback, T
-
-_sim_backends: Dict[Signal, SimSignalBackend] = {}
+from .soft_signal_backend import SoftSignalBackend
+from .utils import DEFAULT_TIMEOUT, Callback, T
 
 
 def _add_timeout(func):
@@ -61,17 +61,19 @@ class Signal(Device, Generic[T]):
         timeout: Optional[float] = DEFAULT_TIMEOUT,
         name: str = "",
     ) -> None:
-        super().__init__(name)
         self._timeout = timeout
-        self._init_backend = self._backend = backend
+        self._initial_backend = self._backend = backend
+        super().__init__(name)
 
-    async def connect(self, sim=False, timeout=DEFAULT_TIMEOUT):
-        if sim:
-            self._backend = SimSignalBackend(datatype=self._init_backend.datatype)
-            _sim_backends[self] = self._backend
-        else:
-            self._backend = self._init_backend
-            _sim_backends.pop(self, None)
+    async def connect(
+        self, mock=False, timeout=DEFAULT_TIMEOUT, force_reconnect: bool = False
+    ):
+        if mock and not isinstance(self._backend, MockSignalBackend):
+            # Using a soft backend, look to the initial value
+            self._backend = MockSignalBackend(
+                initial_backend=self._initial_backend,
+            )
+        self.log.debug(f"Connecting to {self.source}")
         await self._backend.connect(timeout=timeout)
 
     @property
@@ -96,10 +98,12 @@ class _SignalCache(Generic[T]):
         self._value: Optional[T] = None
 
         self.backend = backend
+        signal.log.debug(f"Making subscription on source {signal.source}")
         backend.set_callback(self._callback)
 
     def close(self):
         self.backend.set_callback(None)
+        self._signal.log.debug(f"Closing subscription on source {self._signal.source}")
 
     async def get_reading(self) -> Reading:
         await self._valid.wait()
@@ -112,6 +116,10 @@ class _SignalCache(Generic[T]):
         return self._value
 
     def _callback(self, reading: Reading, value: T):
+        self._signal.log.debug(
+            f"Updated subscription: reading of source {self._signal.source} changed"
+            f"from {self._reading} to {reading}"
+        )
         self._reading = reading
         self._value = value
         self._valid.set()
@@ -178,7 +186,9 @@ class SignalR(Signal[T], AsyncReadable, AsyncStageable, Subscribable):
     @_add_timeout
     async def get_value(self, cached: Optional[bool] = None) -> T:
         """The current value"""
-        return await self._backend_or_cache(cached).get_value()
+        value = await self._backend_or_cache(cached).get_value()
+        self.log.debug(f"get_value() on source {self.source} returned {value}")
+        return value
 
     def subscribe_value(self, function: Callback[T]):
         """Subscribe to updates in value of a device"""
@@ -213,8 +223,15 @@ class SignalW(Signal[T], Movable):
         """Set the value and return a status saying when it's done"""
         if timeout is USE_DEFAULT_TIMEOUT:
             timeout = self._timeout
-        coro = self._backend.put(value, wait=wait, timeout=timeout)
-        return AsyncStatus(coro)
+
+        async def do_set():
+            self.log.debug(f"Putting value {value} to backend at source {self.source}")
+            await self._backend.put(value, wait=wait, timeout=timeout)
+            self.log.debug(
+                f"Successfully put value {value} to backend at source {self.source}"
+            )
+
+        return AsyncStatus(do_set())
 
 
 class SignalRW(SignalR[T], SignalW[T], Locatable):
@@ -239,47 +256,42 @@ class SignalX(Signal):
         return AsyncStatus(coro)
 
 
-def set_sim_value(signal: Signal[T], value: T):
-    """Set the value of a signal that is in sim mode."""
-    _sim_backends[signal]._set_value(value)
-
-
-def set_sim_put_proceeds(signal: Signal[T], proceeds: bool):
-    """Allow or block a put with wait=True from proceeding"""
-    event = _sim_backends[signal].put_proceeds
-    if proceeds:
-        event.set()
-    else:
-        event.clear()
-
-
-def set_sim_callback(signal: Signal[T], callback: ReadingValueCallback[T]) -> None:
-    """Monitor the value of a signal that is in sim mode"""
-    return _sim_backends[signal].set_callback(callback)
-
-
 def soft_signal_rw(
     datatype: Optional[Type[T]] = None,
     initial_value: Optional[T] = None,
     name: str = "",
 ) -> SignalRW[T]:
-    """Creates a read-writable Signal with a SimSignalBackend"""
-    signal = SignalRW(SimSignalBackend(datatype, initial_value), name=name)
+    """Creates a read-writable Signal with a SoftSignalBackend"""
+    signal = SignalRW(SoftSignalBackend(datatype, initial_value), name=name)
     return signal
 
 
-def soft_signal_r_and_backend(
+def soft_signal_r_and_setter(
     datatype: Optional[Type[T]] = None,
     initial_value: Optional[T] = None,
     name: str = "",
-) -> Tuple[SignalR[T], SimSignalBackend]:
-    """Returns a tuple of a read-only Signal and its SimSignalBackend through
+) -> Tuple[SignalR[T], Callable[[T]]]:
+    """Returns a tuple of a read-only Signal and a callable through
     which the signal can be internally modified within the device. Use
     soft_signal_rw if you want a device that is externally modifiable
     """
-    backend = SimSignalBackend(datatype, initial_value)
+    backend = SoftSignalBackend(datatype, initial_value)
     signal = SignalR(backend, name=name)
-    return (signal, backend)
+
+    return (signal, backend.set_value)
+
+
+def _generate_assert_error_msg(
+    name: str, expected_result: str, actuall_result: str
+) -> str:
+    WARNING = "\033[93m"
+    FAIL = "\033[91m"
+    ENDC = "\033[0m"
+    return (
+        f"Expected {WARNING}{name}{ENDC} to produce"
+        + f"\n{FAIL}{actuall_result}{ENDC}"
+        + f"\nbut actually got \n{FAIL}{expected_result}{ENDC}"
+    )
 
 
 async def assert_value(signal: SignalR[T], value: Any) -> None:
@@ -298,11 +310,14 @@ async def assert_value(signal: SignalR[T], value: Any) -> None:
         await assert_value(signal, value)
 
     """
-    assert await signal.get_value() == value
+    actual_value = await signal.get_value()
+    assert actual_value == value, _generate_assert_error_msg(
+        signal.name, value, actual_value
+    )
 
 
 async def assert_reading(
-    readable: AsyncReadable, reading: Mapping[str, Reading]
+    readable: AsyncReadable, expected_reading: Mapping[str, Reading]
 ) -> None:
     """Assert readings from readable.
 
@@ -320,7 +335,10 @@ async def assert_reading(
         await assert_reading(readable, reading)
 
     """
-    assert await readable.read() == reading
+    actual_reading = await readable.read()
+    assert expected_reading == actual_reading, _generate_assert_error_msg(
+        readable.name, expected_reading, actual_reading
+    )
 
 
 async def assert_configuration(
@@ -343,7 +361,10 @@ async def assert_configuration(
         await assert_configuration(configurable configuration)
 
     """
-    assert await configurable.read_configuration() == configuration
+    actual_configurable = await configurable.read_configuration()
+    assert configuration == actual_configurable, _generate_assert_error_msg(
+        configurable.name, configuration, actual_configurable
+    )
 
 
 def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int):
@@ -363,11 +384,18 @@ def assert_emitted(docs: Mapping[str, list[dict]], **numbers: int):
         assert_emitted(docs, start=1, descriptor=1,
         resource=1, datum=1, event=1, stop=1)
     """
-    assert list(docs) == list(numbers)
-    assert {name: len(d) for name, d in docs.items()} == numbers
-
-
-async def observe_value(signal: SignalR[T], timeout=None) -> AsyncGenerator[T, None]:
+    assert list(docs) == list(numbers), _generate_assert_error_msg(
+        "documents", list(numbers), list(docs)
+    )
+    actual_numbers = {name: len(d) for name, d in docs.items()}
+    assert actual_numbers == numbers, _generate_assert_error_msg(
+        "emitted", numbers, actual_numbers
+    )
+
+
+async def observe_value(
+    signal: SignalR[T], timeout=None, done_status: Status | None = None
+) -> AsyncGenerator[T, None]:
     """Subscribe to the value of a signal so it can be iterated from.
 
     Parameters
@@ -375,6 +403,8 @@ async def observe_value(signal: SignalR[T], timeout=None) -> AsyncGenerator[T, N
     signal:
         Call subscribe_value on this at the start, and clear_sub on it at the
         end
+    done_status:
+        If this status is complete, stop observing and make the iterator return.
 
     Notes
     -----
@@ -383,7 +413,10 @@ async def observe_value(signal: SignalR[T], timeout=None) -> AsyncGenerator[T, N
         async for value in observe_value(sig):
             do_something_with(value)
     """
-    q: asyncio.Queue[T] = asyncio.Queue()
+
+    class StatusIsDone: ...
+
+    q: asyncio.Queue[T | StatusIsDone] = asyncio.Queue()
     if timeout is None:
         get_value = q.get
     else:
@@ -391,10 +424,17 @@ async def observe_value(signal: SignalR[T], timeout=None) -> AsyncGenerator[T, N
         async def get_value():
             return await asyncio.wait_for(q.get(), timeout)
 
+    if done_status is not None:
+        done_status.add_callback(lambda _: q.put_nowait(StatusIsDone()))
+
     signal.subscribe_value(q.put_nowait)
     try:
         while True:
-            yield await get_value()
+            item = await get_value()
+            if not isinstance(item, StatusIsDone):
+                yield item
+            else:
+                break
     finally:
         signal.clear_sub(q.put_nowait)
 

ophyd_async/core/{sim_signal_backend.py → soft_signal_backend.py}

@@ -1,12 +1,11 @@
 from __future__ import annotations
 
-import asyncio
 import inspect
 import time
 from collections import abc
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, Generic, Optional, Type, Union, cast, get_origin
+from typing import Dict, Generic, Optional, Type, Union, cast, get_origin
 
 import numpy as np
 from bluesky.protocols import DataKey, Dtype, Reading
@@ -22,7 +21,7 @@ primitive_dtypes: Dict[type, Dtype] = {
 }
 
 
-class SimConverter(Generic[T]):
+class SoftConverter(Generic[T]):
     def value(self, value: T) -> T:
         return value
 
@@ -55,7 +54,7 @@ class SimConverter(Generic[T]):
         return datatype()
 
 
-class SimArrayConverter(SimConverter):
+class SoftArrayConverter(SoftConverter):
     def get_datakey(self, source: str, value) -> DataKey:
         return {"source": source, "dtype": "array", "shape": [len(value)]}
 
@@ -70,7 +69,7 @@ class SimArrayConverter(SimConverter):
 
 
 @dataclass
-class SimEnumConverter(SimConverter):
+class SoftEnumConverter(SoftConverter):
     enum_class: Type[Enum]
 
     def write_value(self, value: Union[Enum, str]) -> Enum:
@@ -90,26 +89,21 @@ class SimEnumConverter(SimConverter):
         return cast(T, list(datatype.__members__.values())[0])  # type: ignore
 
 
-class DisconnectedSimConverter(SimConverter):
-    def __getattribute__(self, __name: str) -> Any:
-        raise NotImplementedError("No PV has been set as connect() has not been called")
-
-
 def make_converter(datatype):
     is_array = get_dtype(datatype) is not None
     is_sequence = get_origin(datatype) == abc.Sequence
     is_enum = issubclass(datatype, Enum) if inspect.isclass(datatype) else False
 
     if is_array or is_sequence:
-        return SimArrayConverter()
+        return SoftArrayConverter()
     if is_enum:
-        return SimEnumConverter(datatype)
+        return SoftEnumConverter(datatype)
 
-    return SimConverter()
+    return SoftConverter()
 
 
-class SimSignalBackend(SignalBackend[T]):
-    """An simulated backend to a Signal, created with ``Signal.connect(sim=True)``"""
+class SoftSignalBackend(SignalBackend[T]):
+    """An backend to a soft Signal, for test signals see ``MockSignalBackend``."""
 
     _value: T
     _initial_value: Optional[T]
@@ -122,25 +116,23 @@ class SimSignalBackend(SignalBackend[T]):
         initial_value: Optional[T] = None,
     ) -> None:
         self.datatype = datatype
-        self.converter: SimConverter = DisconnectedSimConverter()
         self._initial_value = initial_value
-        self.put_proceeds = asyncio.Event()
-        self.put_proceeds.set()
-        self.callback: Optional[ReadingValueCallback[T]] = None
-
-    def source(self, name: str) -> str:
-        return f"soft://{name}"
-
-    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
-        self.converter = make_converter(self.datatype)
+        self.converter: SoftConverter = make_converter(datatype)
         if self._initial_value is None:
             self._initial_value = self.converter.make_initial_value(self.datatype)
         else:
-            # convert potentially unconverted initial value passed to init method
             self._initial_value = self.converter.write_value(self._initial_value)
+
+        self.callback: Optional[ReadingValueCallback[T]] = None
         self._severity = 0
+        self.set_value(self._initial_value)
 
-        await self.put(None)
+    def source(self, name: str) -> str:
+        return f"soft://{name}"
+
+    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
+        """Connection isn't required for soft signals."""
+        pass
 
     async def put(self, value: Optional[T], wait=True, timeout=None):
         write_value = (
@@ -148,13 +140,11 @@ class SimSignalBackend(SignalBackend[T]):
             if value is not None
             else self._initial_value
         )
-        self._set_value(write_value)
 
-        if wait:
-            await asyncio.wait_for(self.put_proceeds.wait(), timeout)
+        self.set_value(write_value)
 
-    def _set_value(self, value: T):
-        """Method to bypass asynchronous logic, designed to only be used in tests."""
+    def set_value(self, value: T):
+        """Method to bypass asynchronous logic."""
         self._value = value
         self._timestamp = time.monotonic()
         reading: Reading = self.converter.reading(
@@ -174,7 +164,7 @@ class SimSignalBackend(SignalBackend[T]):
         return self.converter.value(self._value)
 
     async def get_setpoint(self) -> T:
-        """For a simulated backend, the setpoint and readback values are the same."""
+        """For a soft signal, the setpoint and readback values are the same."""
         return await self.get_value()
 
     def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:

ophyd_async/core/utils.py

@@ -2,13 +2,16 @@ from __future__ import annotations
 
 import asyncio
 import logging
+from dataclasses import dataclass
 from typing import (
     Awaitable,
     Callable,
     Dict,
+    Generic,
     Iterable,
     List,
     Optional,
+    ParamSpec,
     Type,
     TypeVar,
     Union,
@@ -18,6 +21,7 @@ import numpy as np
 from bluesky.protocols import Reading
 
 T = TypeVar("T")
+P = ParamSpec("P")
 Callback = Callable[[T], None]
 
 #: A function that will be called with the Reading and value when the
@@ -77,6 +81,21 @@ class NotConnected(Exception):
         return self.format_error_string(indent="")
 
 
+@dataclass(frozen=True)
+class WatcherUpdate(Generic[T]):
+    """A dataclass such that, when expanded, it provides the kwargs for a watcher"""
+
+    current: T
+    initial: T
+    target: T
+    name: str | None = None
+    unit: str | None = None
+    precision: float | None = None
+    fraction: float | None = None
+    time_elapsed: float | None = None
+    time_remaining: float | None = None
+
+
 async def wait_for_connection(**coros: Awaitable[None]):
     """Call many underlying signals, accumulating exceptions and returning them
 

ophyd_async/epics/_backend/_aioca.py

@@ -28,7 +28,7 @@ from ophyd_async.core import (
 )
 from ophyd_async.core.utils import DEFAULT_TIMEOUT, NotConnected
 
-from .common import get_supported_enum_class
+from .common import get_supported_values
 
 dbr_to_dtype: Dict[Dbr, Dtype] = {
     dbr.DBR_STRING: "string",
@@ -79,7 +79,7 @@ class CaArrayConverter(CaConverter):
 
 @dataclass
 class CaEnumConverter(CaConverter):
-    enum_class: Type[Enum]
+    choices: dict[str, str]
 
     def write_value(self, value: Union[Enum, str]):
         if isinstance(value, Enum):
@@ -88,11 +88,15 @@ class CaEnumConverter(CaConverter):
             return value
 
     def value(self, value: AugmentedValue):
-        return self.enum_class(value)
+        return self.choices[value]
 
     def get_datakey(self, source: str, value: AugmentedValue) -> DataKey:
-        choices = [e.value for e in self.enum_class]
-        return {"source": source, "dtype": "string", "shape": [], "choices": choices}
+        return {
+            "source": source,
+            "dtype": "string",
+            "shape": [],
+            "choices": list(self.choices),
+        }
 
 
 class DisconnectedCaConverter(CaConverter):
@@ -138,8 +142,8 @@ def make_converter(
         pv_choices = get_unique(
             {k: tuple(v.enums) for k, v in values.items()}, "choices"
         )
-        enum_class = get_supported_enum_class(pv, datatype, pv_choices)
-        return CaEnumConverter(dbr.DBR_STRING, None, enum_class)
+        supported_values = get_supported_values(pv, datatype, pv_choices)
+        return CaEnumConverter(dbr.DBR_STRING, None, supported_values)
     else:
         value = list(values.values())[0]
         # Done the dbr check, so enough to check one of the values