ophyd-async 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

ophyd_async/core/device_save_loader.py

@@ -0,0 +1,286 @@
+from enum import Enum
+from functools import partial
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Union
+
+import numpy as np
+import numpy.typing as npt
+import yaml
+from bluesky.plan_stubs import abs_set, wait
+from bluesky.protocols import Location
+from bluesky.utils import Msg
+from epicscorelibs.ca.dbr import ca_array, ca_float, ca_int, ca_str
+
+from .device import Device
+from .signal import SignalRW
+
+CaType = Union[ca_float, ca_int, ca_str, ca_array]
+
+
+def ndarray_representer(dumper: yaml.Dumper, array: npt.NDArray[Any]) -> yaml.Node:
+    return dumper.represent_sequence(
+        "tag:yaml.org,2002:seq", array.tolist(), flow_style=True
+    )
+
+
+def ca_dbr_representer(dumper: yaml.Dumper, value: CaType) -> yaml.Node:
+    # if it's an array, just call ndarray_representer...
+    represent_array = partial(ndarray_representer, dumper)
+
+    representers: Dict[CaType, Callable[[CaType], yaml.Node]] = {
+        ca_float: dumper.represent_float,
+        ca_int: dumper.represent_int,
+        ca_str: dumper.represent_str,
+        ca_array: represent_array,
+    }
+    return representers[type(value)](value)
+
+
+class OphydDumper(yaml.Dumper):
+    def represent_data(self, data: Any) -> Any:
+        if isinstance(data, Enum):
+            return self.represent_data(data.value)
+        return super(OphydDumper, self).represent_data(data)
+
+
+def get_signal_values(
+    signals: Dict[str, SignalRW[Any]], ignore: Optional[List[str]] = None
+) -> Generator[Msg, Sequence[Location[Any]], Dict[str, Any]]:
+    """Get signal values in bulk.
+
+    Used as part of saving the signals of a device to a yaml file.
+
+    Parameters
+    ----------
+    signals : Dict[str, SignalRW]
+        Dictionary with pv names and matching SignalRW values. Often the direct result
+        of :func:`walk_rw_signals`.
+
+    ignore : Optional[List[str]]
+        Optional list of PVs that should be ignored.
+
+    Returns
+    -------
+    Dict[str, Any]
+        A dictionary containing pv names and their associated values. Ignored pvs are
+        set to None.
+
+    See Also
+    --------
+    :func:`ophyd_async.core.walk_rw_signals`
+    :func:`ophyd_async.core.save_to_yaml`
+    """
+
+    ignore = ignore or []
+    selected_signals = {
+        key: signal for key, signal in signals.items() if key not in ignore
+    }
+    selected_values = yield Msg("locate", *selected_signals.values())
+
+    # TODO: investigate wrong type hints
+    if isinstance(selected_values, dict):
+        selected_values = [selected_values]  # type: ignore
+
+    assert selected_values is not None, "No signalRW's were able to be located"
+    named_values = {
+        key: value["setpoint"] for key, value in zip(selected_signals, selected_values)
+    }
+    # Ignored values place in with value None so we know which ones were ignored
+    named_values.update({key: None for key in ignore})
+    return named_values
+
+
+def walk_rw_signals(
+    device: Device, path_prefix: Optional[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.
+
+        See Also
+    --------
+    :func:`ophyd_async.core.get_signal_values`
+    :func:`ophyd_async.core.save_to_yaml`
+
+    """
+
+    if not path_prefix:
+        path_prefix = ""
+
+    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
+
+
+def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
+    """Plan which serialises a phase or set of phases of SignalRWs to a yaml file.
+
+    Parameters
+    ----------
+    phases : dict or list of dicts
+        The values to save. Each item in the list is a seperate phase used when loading
+        a device. In general this variable be the return value of `get_signal_values`.
+
+    save_path : str
+        Path of the yaml file to write to
+
+    See Also
+    --------
+    :func:`ophyd_async.core.walk_rw_signals`
+    :func:`ophyd_async.core.get_signal_values`
+    :func:`ophyd_async.core.load_from_yaml`
+    """
+
+    yaml.add_representer(np.ndarray, ndarray_representer, Dumper=yaml.Dumper)
+
+    yaml.add_representer(ca_float, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_int, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_str, ca_dbr_representer, Dumper=yaml.Dumper)
+    yaml.add_representer(ca_array, ca_dbr_representer, Dumper=yaml.Dumper)
+
+    with open(save_path, "w") as file:
+        yaml.dump(phases, file, Dumper=OphydDumper, default_flow_style=False)
+
+
+def load_from_yaml(save_path: str) -> Sequence[Dict[str, Any]]:
+    """Plan that returns a list of dicts with saved signal values from a yaml file.
+
+    Parameters
+    ----------
+    save_path : str
+        Path of the yaml file to load from
+
+    See Also
+    --------
+    :func:`ophyd_async.core.save_to_yaml`
+    :func:`ophyd_async.core.set_signal_values`
+    """
+    with open(save_path, "r") as file:
+        return yaml.full_load(file)
+
+
+def set_signal_values(
+    signals: Dict[str, SignalRW[Any]], values: Sequence[Dict[str, Any]]
+) -> Generator[Msg, None, None]:
+    """Maps signals from a yaml file into device signals.
+
+    ``values`` contains signal values in phases, which are loaded in sequentially
+    into the provided signals, to ensure signals are set in the correct order.
+
+    Parameters
+    ----------
+    signals : Dict[str, SignalRW[Any]]
+        Dictionary of named signals to be updated if value found in values argument.
+        Can be the output of :func:`walk_rw_signals()` for a device.
+
+    values : Sequence[Dict[str, Any]]
+        List of dictionaries of signal name and value pairs, if a signal matches
+        the name of one in the signals argument, sets the signal to that value.
+        The groups of signals are loaded in their list order.
+        Can be the output of :func:`load_from_yaml()` for a yaml file.
+
+    See Also
+    --------
+    :func:`ophyd_async.core.load_from_yaml`
+    :func:`ophyd_async.core.walk_rw_signals`
+    """
+    # For each phase, set all the signals,
+    # load them to the correct value and wait for the load to complete
+    for phase_number, phase in enumerate(values):
+        # Key is signal name
+        for key, value in phase.items():
+            # Skip ignored values
+            if value is None:
+                continue
+
+            if key in signals:
+                yield from abs_set(
+                    signals[key], value, group=f"load-phase{phase_number}"
+                )
+
+        yield from wait(f"load-phase{phase_number}")
+
+
+def load_device(device: Device, path: str):
+    """Plan which loads PVs from a yaml file into a device.
+
+    Parameters
+    ----------
+    device: Device
+        The device to load PVs into
+    path: str
+        Path of the yaml file to load
+
+    See Also
+    --------
+    :func:`ophyd_async.core.save_device`
+    """
+    values = load_from_yaml(path)
+    signals_to_set = walk_rw_signals(device)
+    yield from set_signal_values(signals_to_set, values)
+
+
+def all_at_once(values: Dict[str, Any]) -> Sequence[Dict[str, Any]]:
+    """Sort all the values into a single phase so they are set all at once"""
+    return [values]
+
+
+def save_device(
+    device: Device,
+    path: str,
+    sorter: Callable[[Dict[str, Any]], Sequence[Dict[str, Any]]] = all_at_once,
+    ignore: Optional[List[str]] = None,
+):
+    """Plan that saves the state of all PV's on a device using a sorter.
+
+    The default sorter assumes all saved PVs can be loaded at once, and therefore
+    can be saved at one time, i.e. all PVs will appear on one list in the
+    resulting yaml file.
+
+    This can be a problem, because when the yaml is ingested with
+    :func:`ophyd_async.core.load_device`, it will set all of those PVs at once.
+    However, some PV's need to be set before others - this is device specific.
+
+    Therefore, users should consider the order of device loading and write their
+    own sorter algorithms accordingly.
+
+    See :func:`ophyd_async.panda.phase_sorter` for a valid implementation of the
+    sorter.
+
+    Parameters
+    ----------
+    device : Device
+        The device whose PVs should be saved.
+
+    path : str
+        The path where the resulting yaml should be saved to
+
+    sorter : Callable[[Dict[str, Any]], Sequence[Dict[str, Any]]]
+
+    ignore : Optional[List[str]]
+
+    See Also
+    --------
+    :func:`ophyd_async.core.load_device`
+    """
+    values = yield from get_signal_values(walk_rw_signals(device), ignore=ignore)
+    save_to_yaml(sorter(values), path)

ophyd_async/core/flyer.py

@@ -0,0 +1,85 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Generic, Sequence, TypeVar
+
+from bluesky.protocols import DataKey, Flyable, Preparable, Reading, Stageable
+
+from .async_status import AsyncStatus
+from .device import Device
+from .signal import SignalR
+from .utils import merge_gathered_dicts
+
+T = TypeVar("T")
+
+
+class TriggerLogic(ABC, Generic[T]):
+    @abstractmethod
+    async def prepare(self, value: T):
+        """Move to the start of the flyscan"""
+
+    @abstractmethod
+    async def kickoff(self):
+        """Start the flyscan"""
+
+    @abstractmethod
+    async def complete(self):
+        """Block until the flyscan is done"""
+
+    @abstractmethod
+    async def stop(self):
+        """Stop flying and wait everything to be stopped"""
+
+
+class HardwareTriggeredFlyable(
+    Device,
+    Stageable,
+    Preparable,
+    Flyable,
+    Generic[T],
+):
+    def __init__(
+        self,
+        trigger_logic: TriggerLogic[T],
+        configuration_signals: Sequence[SignalR] = (),
+        name: str = "",
+    ):
+        self._trigger_logic = trigger_logic
+        self._configuration_signals = tuple(configuration_signals)
+        super().__init__(name=name)
+
+    @property
+    def trigger_logic(self) -> TriggerLogic[T]:
+        return self._trigger_logic
+
+    @AsyncStatus.wrap
+    async def stage(self) -> None:
+        await self.unstage()
+
+    @AsyncStatus.wrap
+    async def unstage(self) -> None:
+        await self._trigger_logic.stop()
+
+    def prepare(self, value: T) -> AsyncStatus:
+        """Setup trajectories"""
+        return AsyncStatus(self._prepare(value))
+
+    async def _prepare(self, value: T) -> None:
+        # Move to start and setup the flyscan
+        await self._trigger_logic.prepare(value)
+
+    @AsyncStatus.wrap
+    async def kickoff(self) -> None:
+        await self._trigger_logic.kickoff()
+
+    @AsyncStatus.wrap
+    async def complete(self) -> None:
+        await self._trigger_logic.complete()
+
+    async def describe_configuration(self) -> Dict[str, DataKey]:
+        return await merge_gathered_dicts(
+            [sig.describe() for sig in self._configuration_signals]
+        )
+
+    async def read_configuration(self) -> Dict[str, Reading]:
+        return await merge_gathered_dicts(
+            [sig.read() for sig in self._configuration_signals]
+        )

ophyd_async/core/mock_signal_backend.py

@@ -0,0 +1,82 @@
+import asyncio
+from functools import cached_property
+from typing import Callable, Optional, Type
+from unittest.mock import Mock
+
+from bluesky.protocols import Descriptor, Reading
+
+from ophyd_async.core.signal_backend import SignalBackend
+from ophyd_async.core.soft_signal_backend import SoftSignalBackend
+from ophyd_async.core.utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
+
+
+class MockSignalBackend(SignalBackend[T]):
+    def __init__(
+        self,
+        datatype: Optional[Type[T]] = None,
+        initial_backend: Optional[SignalBackend[T]] = None,
+    ) -> None:
+        if isinstance(initial_backend, MockSignalBackend):
+            raise ValueError("Cannot make a MockSignalBackend for a MockSignalBackends")
+
+        self.initial_backend = initial_backend
+
+        if datatype is None:
+            assert (
+                self.initial_backend
+            ), "Must supply either initial_backend or datatype"
+            datatype = self.initial_backend.datatype
+
+        self.datatype = datatype
+
+        if not isinstance(self.initial_backend, SoftSignalBackend):
+            # If the backend is a hard signal backend, or not provided,
+            # then we create a soft signal to mimic it
+
+            self.soft_backend = SoftSignalBackend(datatype=datatype)
+        else:
+            self.soft_backend = self.initial_backend
+
+    def source(self, name: str) -> str:
+        if self.initial_backend:
+            return f"mock+{self.initial_backend.source(name)}"
+        return f"mock+{name}"
+
+    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
+        pass
+
+    @cached_property
+    def put_mock(self) -> Mock:
+        return Mock(name="put", spec=Callable)
+
+    @cached_property
+    def put_proceeds(self) -> asyncio.Event:
+        put_proceeds = asyncio.Event()
+        put_proceeds.set()
+        return put_proceeds
+
+    async def put(self, value: Optional[T], wait=True, timeout=None):
+        self.put_mock(value, wait=wait, timeout=timeout)
+        await self.soft_backend.put(value, wait=wait, timeout=timeout)
+
+        if wait:
+            await asyncio.wait_for(self.put_proceeds.wait(), timeout=timeout)
+
+    def set_value(self, value: T):
+        self.soft_backend.set_value(value)
+
+    async def get_reading(self) -> Reading:
+        return await self.soft_backend.get_reading()
+
+    async def get_value(self) -> T:
+        return await self.soft_backend.get_value()
+
+    async def get_setpoint(self) -> T:
+        """For a soft signal, the setpoint and readback values are the same."""
+        return await self.soft_backend.get_setpoint()
+
+    async def get_datakey(self, source: str) -> Descriptor:
+        return await self.soft_backend.get_datakey(source)
+
+    def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
+        self.soft_backend.set_callback(callback)

ophyd_async/core/mock_signal_utils.py

@@ -0,0 +1,145 @@
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, Callable, Iterable
+from unittest.mock import 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, 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: Signal):
+    for signal in signals:
+        set_mock_put_proceeds(signal, False)
+    yield
+    for signal in signals:
+        set_mock_put_proceeds(signal, True)
+
+
+def get_mock_put(signal: Signal) -> Mock:
+    """Get the mock associated with the put call on the signal."""
+    return _get_mock_signal_backend(signal).put_mock
+
+
+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(list(self.values)):
+            raise AssertionError("Not all values have been consumed.")
+
+
+def set_mock_values(
+    signal: Signal,
+    values: Iterable[Any],
+    require_all_consumed: bool = False,
+) -> _SetValuesIterator:
+    """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
+
+
+def callback_on_mock_put(signal: Signal[T], callback: Callable[[T], 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)