ophyd-async 0.1.0__py3-none-any.whl → 0.3a1__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,94 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Generic, Optional, Sequence, TypeVar
+
+from bluesky.protocols import Descriptor, Flyable, Preparable, Reading, Stageable
+
+from .async_status import AsyncStatus
+from .detector import TriggerInfo
+from .device import Device
+from .signal import SignalR
+from .utils import merge_gathered_dicts
+
+T = TypeVar("T")
+
+
+class TriggerLogic(ABC, Generic[T]):
+    @abstractmethod
+    def trigger_info(self, value: T) -> TriggerInfo:
+        """Return info about triggers that will be produced for a given value"""
+
+    @abstractmethod
+    async def prepare(self, value: T):
+        """Move to the start of the flyscan"""
+
+    @abstractmethod
+    async def start(self):
+        """Start the flyscan"""
+
+    @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)
+        self._describe: Dict[str, Descriptor] = {}
+        self._fly_status: Optional[AsyncStatus] = None
+        self._trigger_info: Optional[TriggerInfo] = None
+        super().__init__(name=name)
+
+    @property
+    def trigger_logic(self) -> TriggerLogic[T]:
+        return self._trigger_logic
+
+    @property
+    def trigger_info(self) -> Optional[TriggerInfo]:
+        return self._trigger_info
+
+    @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:
+        self._trigger_info = self._trigger_logic.trigger_info(value)
+        # Move to start and setup the flyscan
+        await self._trigger_logic.prepare(value)
+
+    @AsyncStatus.wrap
+    async def kickoff(self) -> None:
+        self._fly_status = AsyncStatus(self._trigger_logic.start())
+
+    def complete(self) -> AsyncStatus:
+        assert self._fly_status, "Kickoff not run"
+        return self._fly_status
+
+    async def describe_configuration(self) -> Dict[str, Descriptor]:
+        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/{_device/_signal/signal.py → signal.py}

@@ -6,6 +6,8 @@ from typing import AsyncGenerator, Callable, Dict, Generic, Optional, Union
 
 from bluesky.protocols import (
     Descriptor,
+    Locatable,
+    Location,
     Movable,
     Readable,
     Reading,
@@ -13,11 +15,11 @@ from bluesky.protocols import (
     Subscribable,
 )
 
-from ...async_status import AsyncStatus
-from ...utils import DEFAULT_TIMEOUT, Callback, ReadingValueCallback, T
-from .._backend.signal_backend import SignalBackend
-from .._backend.sim_signal_backend import SimSignalBackend
-from ..device import Device
+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] = {}
 
@@ -56,7 +58,7 @@ class Signal(Device, Generic[T]):
     def set_name(self, name: str = ""):
         self._name = name
 
-    async def connect(self, sim=False):
+    async def connect(self, sim=False, timeout=DEFAULT_TIMEOUT):
         if sim:
             self._backend = SimSignalBackend(
                 datatype=self._init_backend.datatype, source=self._init_backend.source
@@ -65,7 +67,7 @@ class Signal(Device, Generic[T]):
         else:
             self._backend = self._init_backend
             _sim_backends.pop(self, None)
-        await self._backend.connect()
+        await self._backend.connect(timeout=timeout)
 
     @property
     def source(self) -> str:
@@ -196,25 +198,40 @@ class SignalR(Signal[T], Readable, Stageable, Subscribable):
         self._del_cache(self._get_cache().set_staged(False))
 
 
+USE_DEFAULT_TIMEOUT = "USE_DEFAULT_TIMEOUT"
+
+
 class SignalW(Signal[T], Movable):
     """Signal that can be set"""
 
-    def set(self, value: T, wait=True, timeout=None) -> AsyncStatus:
+    def set(self, value: T, wait=True, timeout=USE_DEFAULT_TIMEOUT) -> AsyncStatus:
         """Set the value and return a status saying when it's done"""
-        coro = self._backend.put(value, wait=wait, timeout=timeout or self._timeout)
+        if timeout is USE_DEFAULT_TIMEOUT:
+            timeout = self._timeout
+        coro = self._backend.put(value, wait=wait, timeout=timeout)
         return AsyncStatus(coro)
 
 
-class SignalRW(SignalR[T], SignalW[T]):
+class SignalRW(SignalR[T], SignalW[T], Locatable):
     """Signal that can be both read and set"""
 
+    async def locate(self) -> Location:
+        location: Location = {
+            "setpoint": await self._backend.get_setpoint(),
+            "readback": await self.get_value(),
+        }
+        return location
+
 
 class SignalX(Signal):
     """Signal that puts the default value"""
 
-    async def execute(self, wait=True, timeout=None):
-        """Execute the action and return a status saying when it's done"""
-        await self._backend.put(None, wait=wait, timeout=timeout or self._timeout)
+    def trigger(self, wait=True, timeout=USE_DEFAULT_TIMEOUT) -> AsyncStatus:
+        """Trigger the action and return a status saying when it's done"""
+        if timeout is USE_DEFAULT_TIMEOUT:
+            timeout = self._timeout
+        coro = self._backend.put(None, wait=wait, timeout=timeout)
+        return AsyncStatus(coro)
 
 
 def set_sim_value(signal: Signal[T], value: T):
@@ -236,7 +253,7 @@ def set_sim_callback(signal: Signal[T], callback: ReadingValueCallback[T]) -> No
     return _sim_backends[signal].set_callback(callback)
 
 
-async def observe_value(signal: SignalR[T]) -> AsyncGenerator[T, None]:
+async def observe_value(signal: SignalR[T], timeout=None) -> AsyncGenerator[T, None]:
     """Subscribe to the value of a signal so it can be iterated from.
 
     Parameters
@@ -253,17 +270,24 @@ async def observe_value(signal: SignalR[T]) -> AsyncGenerator[T, None]:
             do_something_with(value)
     """
     q: asyncio.Queue[T] = asyncio.Queue()
+    if timeout is None:
+        get_value = q.get
+    else:
+
+        async def get_value():
+            return await asyncio.wait_for(q.get(), timeout)
+
     signal.subscribe_value(q.put_nowait)
     try:
         while True:
-            yield await q.get()
+            yield await get_value()
     finally:
         signal.clear_sub(q.put_nowait)
 
 
 class _ValueChecker(Generic[T]):
     def __init__(self, matcher: Callable[[T], bool], matcher_name: str):
-        self._last_value: Optional[T]
+        self._last_value: Optional[T] = None
         self._matcher = matcher
         self._matcher_name = matcher_name
 
@@ -273,7 +297,7 @@ class _ValueChecker(Generic[T]):
             if self._matcher(value):
                 return
 
-    async def wait_for_value(self, signal: SignalR[T], timeout: float):
+    async def wait_for_value(self, signal: SignalR[T], timeout: Optional[float]):
         try:
             await asyncio.wait_for(self._wait_for_value(signal), timeout)
         except asyncio.TimeoutError as e:
@@ -284,7 +308,7 @@ class _ValueChecker(Generic[T]):
 
 
 async def wait_for_value(
-    signal: SignalR[T], match: Union[T, Callable[[T], bool]], timeout: float
+    signal: SignalR[T], match: Union[T, Callable[[T], bool]], timeout: Optional[float]
 ):
     """Wait for a signal to have a matching value.
 
@@ -330,6 +354,10 @@ async def set_and_wait_for_value(
     - Read the same Signal to check the operation has started
     - Return the Status so calling code can wait for operation to complete
 
+    This function sets a signal to a specified value, optionally with or without a
+    ca/pv put callback, and waits for the readback value of the signal to match the
+    value it was set to.
+
     Parameters
     ----------
     signal:

ophyd_async/core/{_device/_backend/signal_backend.py → signal_backend.py}

@@ -3,7 +3,7 @@ from typing import Generic, Optional, Type
 
 from bluesky.protocols import Descriptor, Reading
 
-from ...utils import ReadingValueCallback, T
+from .utils import DEFAULT_TIMEOUT, ReadingValueCallback, T
 
 
 class SignalBackend(Generic[T]):
@@ -16,7 +16,7 @@ class SignalBackend(Generic[T]):
     source: str = ""
 
     @abstractmethod
-    async def connect(self):
+    async def connect(self, timeout: float = DEFAULT_TIMEOUT):
         """Connect to underlying hardware"""
 
     @abstractmethod
@@ -35,6 +35,10 @@ class SignalBackend(Generic[T]):
     async def get_value(self) -> T:
         """The current value"""
 
+    @abstractmethod
+    async def get_setpoint(self) -> T:
+        """The point that a signal was requested to move to."""
+
     @abstractmethod
     def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
         """Observe changes to the current value, timestamp and severity"""

ophyd_async/core/{_device/_backend/sim_signal_backend.py → sim_signal_backend.py}

@@ -11,8 +11,8 @@ from typing import Any, Dict, Generic, Optional, Type, Union, cast, get_origin
 
 from bluesky.protocols import Descriptor, Dtype, Reading
 
-from ...utils import ReadingValueCallback, T, get_dtype
 from .signal_backend import SignalBackend
+from .utils import DEFAULT_TIMEOUT, ReadingValueCallback, T, get_dtype
 
 primitive_dtypes: Dict[type, Dtype] = {
     str: "string",
@@ -123,7 +123,7 @@ class SimSignalBackend(SignalBackend[T]):
         self.put_proceeds.set()
         self.callback: Optional[ReadingValueCallback[T]] = None
 
-    async def connect(self) -> None:
+    async def connect(self, timeout: float = DEFAULT_TIMEOUT) -> None:
         self.converter = make_converter(self.datatype)
         self._initial_value = self.converter.make_initial_value(self.datatype)
         self._severity = 0
@@ -161,6 +161,10 @@ class SimSignalBackend(SignalBackend[T]):
     async def get_value(self) -> 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."""
+        return await self.get_value()
+
     def set_callback(self, callback: Optional[ReadingValueCallback[T]]) -> None:
         if callback:
             assert not self.callback, "Cannot set a callback when one is already set"

ophyd_async/core/{_device/standard_readable.py → standard_readable.py}

@@ -2,10 +2,10 @@ from typing import Dict, Sequence, Tuple
 
 from bluesky.protocols import Configurable, Descriptor, Readable, Reading, Stageable
 
-from ..async_status import AsyncStatus
-from ..utils import merge_gathered_dicts
-from ._signal.signal import SignalR
+from .async_status import AsyncStatus
 from .device import Device
+from .signal import SignalR
+from .utils import merge_gathered_dicts
 
 
 class StandardReadable(Device, Readable, Configurable, Stageable):