Add panda-specific save load plans

src/ophyd_async/core/__init__.py

@@ -10,7 +10,9 @@
 from .device import Device, DeviceCollector, DeviceVector
 from .device_save_loader import (
     get_signal_values,
+    load_device,
     load_from_yaml,
+    save_device,
     save_to_yaml,
     set_signal_values,
     walk_rw_signals,
@@ -97,4 +99,6 @@
     "save_to_yaml",
     "set_signal_values",
     "walk_rw_signals",
+    "load_device",
+    "save_device",
 ]

src/ophyd_async/core/device_save_loader.py

@@ -45,30 +45,29 @@ def represent_data(self, data: Any) -> Any:
 def get_signal_values(
     signals: Dict[str, SignalRW[Any]], ignore: Optional[List[str]] = None
 ) -> Generator[Msg, Sequence[Location[Any]], Dict[str, Any]]:
-    """
-    Read the values of SignalRW's, to be used alongside `walk_rw_signals`. Used as part
-    of saving a device.
+    """Get signal values in bulk.
+
+    Used as part of saving the signals of a device to a yaml file.
+
     Parameters
     ----------
-        signals : Dict[str, SignalRW]: A dictionary matching the string attribute path
-        of a SignalRW to the signal itself
+    signals : Dict[str, SignalRW]
+        Dictionary with pv names and matching SignalRW values. Often the direct result
+        of :func:`walk_rw_signals`.
 
-        ignore : List of strings: . A list of string attribute paths to the SignalRW's
-        to be ignored.
+    ignore : Optional[List[str]]
+        Optional list of PVs that should be ignored.
 
     Returns
     -------
-        Dict[str, Any]: A dictionary matching the string attribute path of a SignalRW
-        to the value of that signal. Ignored attributes are set to None.
-
-    Yields:
-        Iterator[Dict[str, Any]]: The Location of a signal
+    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 []
@@ -93,9 +92,11 @@ def get_signal_values(
 def walk_rw_signals(
     device: Device, path_prefix: Optional[str] = ""
 ) -> Dict[str, SignalRW[Any]]:
-    """
-    Get all the SignalRWs from a device and store them with their dotted attribute
-    paths in a dictionary. Used as part of saving and loading a device
+    """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
@@ -131,9 +132,7 @@ def walk_rw_signals(
 
 
 def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
-    """
-    Serialises and saves a phase or a set of phases of a device's SignalRW's to a
-    yaml file.
+    """Plan which serialises a phase or set of phases of SignalRWs to a yaml file.
 
     Parameters
     ----------
@@ -163,8 +162,7 @@ def save_to_yaml(phases: Sequence[Dict[str, Any]], save_path: str) -> None:
 
 
 def load_from_yaml(save_path: str) -> Sequence[Dict[str, Any]]:
-    """
-    Returns a list of dicts with saved signal values from a yaml file
+    """Plan that returns a list of dicts with saved signal values from a yaml file.
 
     Parameters
     ----------
@@ -183,18 +181,22 @@ def load_from_yaml(save_path: str) -> Sequence[Dict[str, Any]]:
 def set_signal_values(
     signals: Dict[str, SignalRW[Any]], values: Sequence[Dict[str, Any]]
 ) -> Generator[Msg, None, None]:
-    """
-    Loads a phase or a set of phases from a yaml file and puts value to an Ophyd device
+    """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
     --------
@@ -206,8 +208,79 @@ def set_signal_values(
     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)

src/ophyd_async/panda/__init__.py

@@ -6,6 +6,7 @@
     seq_table_from_arrays,
     seq_table_from_rows,
 )
+from .utils import phase_sorter
 
 __all__ = [
     "PandA",
@@ -19,4 +20,5 @@
     "SeqTableRow",
     "SeqTrigger",
     "pvi",
+    "phase_sorter",
 ]

src/ophyd_async/panda/utils.py

@@ -0,0 +1,15 @@
+from typing import Any, Dict, Sequence
+
+
+def phase_sorter(panda_signal_values: Dict[str, Any]) -> Sequence[Dict[str, Any]]:
+    # Panda has two load phases. If the signal name ends in the string "UNITS",
+    # it needs to be loaded first so put in first phase
+    phase_1, phase_2 = {}, {}
+
+    for key, value in panda_signal_values.items():
+        if key.endswith("units"):
+            phase_1[key] = value
+        else:
+            phase_2[key] = value
+
+    return [phase_1, phase_2]

tests/conftest.py

@@ -68,7 +68,10 @@ def pva():
         assert not p.poll(), p.stdout.read()
 
     yield processes
-    [p.terminate() for p in processes]
+
+    for p in processes:
+        p.terminate()
+        p.communicate()
 
 
 @pytest.fixture

tests/core/test_device_save_loader.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from os import path
 from typing import Any, Dict, List
+from unittest.mock import patch
 
 import numpy as np
 import numpy.typing as npt
@@ -13,11 +14,14 @@
     SignalR,
     SignalRW,
     get_signal_values,
+    load_device,
     load_from_yaml,
+    save_device,
     save_to_yaml,
     set_signal_values,
     walk_rw_signals,
 )
+from ophyd_async.core.device_save_loader import all_at_once
 from ophyd_async.epics.signal import epics_signal_r, epics_signal_rw
 
 
@@ -44,13 +48,6 @@ def __init__(self, name: str):
         self.position: npt.NDArray[np.int32]
 
 
-def save_device(device, file_path):
-    signalRWs = walk_rw_signals(device)
-    values = yield from get_signal_values(signalRWs)
-    phases = sort_signal_by_phase(values)
-    save_to_yaml(phases, file_path)
-
-
 @pytest.fixture
 async def device() -> DummyDeviceGroup:
     device = DummyDeviceGroup("parent")
@@ -130,7 +127,7 @@ async def test_yaml_formatting(RE: RunEngine, device, tmp_path):
     await device.child1.sig1.set("test_string")
     table_pv = {"VAL1": np.array([1, 2, 3, 4, 5]), "VAL2": np.array([6, 7, 8, 9, 10])}
     await device.child2.sig1.set(table_pv)
-    RE(save_device(device, file_path))
+    RE(save_device(device, file_path, sorter=sort_signal_by_phase))
 
     with open(file_path, "r") as file:
         expected = """\
@@ -148,7 +145,8 @@ async def test_load_from_yaml(RE: RunEngine, device, tmp_path):
     array = np.array([1, 1, 1, 1, 1])
     await device.child1.sig1.set("initial_string")
     await device.child2.sig1.set(array)
-    RE(save_device(device, file_path))
+    await device.parent_sig1.set(None)
+    RE(save_device(device, file_path, sorter=sort_signal_by_phase))
 
     values = load_from_yaml(file_path)
     assert values[0]["child1.sig1"] == "initial_string"
@@ -160,7 +158,7 @@ async def test_set_signal_values_restores_value(RE: RunEngine, device, tmp_path)
 
     await device.child1.sig1.set("initial_string")
     await device.child2.sig1.set(np.array([1, 1, 1, 1, 1]))
-    RE(save_device(device, file_path))
+    RE(save_device(device, file_path, sorter=sort_signal_by_phase))
 
     await device.child1.sig1.set("changed_string")
     await device.child2.sig1.set(np.array([2, 2, 2, 2, 2]))
@@ -178,3 +176,37 @@ async def test_set_signal_values_restores_value(RE: RunEngine, device, tmp_path)
     array_value = await device.child2.sig1.get_value()
     assert string_value == "initial_string"
     assert np.array_equal(array_value, np.array([1, 1, 1, 1, 1]))
+
+
+@patch("ophyd_async.core.device_save_loader.load_from_yaml")
+@patch("ophyd_async.core.device_save_loader.walk_rw_signals")
+@patch("ophyd_async.core.device_save_loader.set_signal_values")
+async def test_load_device(
+    mock_set_signal_values, mock_walk_rw_signals, mock_load_from_yaml, device
+):
+    RE = RunEngine()
+    RE(load_device(device, "path"))
+    mock_load_from_yaml.assert_called_once()
+    mock_walk_rw_signals.assert_called_once()
+    mock_set_signal_values.assert_called_once()
+
+
+async def test_set_signal_values_skips_ignored_values(device):
+    RE = RunEngine()
+    array = np.array([1, 1, 1, 1, 1])
+
+    await device.child1.sig1.set("initial_string")
+    await device.child2.sig1.set(array)
+    await device.parent_sig1.set(None)
+
+    signals_of_device = walk_rw_signals(device)
+    values_to_set = [{"child1.sig1": None, "child2.sig1": np.array([2, 3, 4])}]
+
+    RE(set_signal_values(signals_of_device, values_to_set))
+
+    assert np.all(await device.child2.sig1.get_value() == np.array([2, 3, 4]))
+    assert await device.child1.sig1.get_value() == "initial_string"
+
+
+def test_all_at_once_sorter():
+    assert all_at_once({"child1.sig1": 0}) == [{"child1.sig1": 0}]

tests/panda/test_panda_utils.py

@@ -0,0 +1,37 @@
+from unittest.mock import patch
+
+import pytest
+from bluesky import RunEngine
+
+from ophyd_async.core import SignalRW, save_device
+from ophyd_async.core.device import DeviceCollector
+from ophyd_async.epics.signal import epics_signal_rw
+from ophyd_async.panda import PandA
+from ophyd_async.panda.utils import phase_sorter
+
+
+@pytest.fixture
+async def sim_panda():
+    async with DeviceCollector(sim=True):
+        sim_panda = PandA("PANDA")
+        sim_panda.phase_1_signal_units: SignalRW = epics_signal_rw(int, "")
+    assert sim_panda.name == "sim_panda"
+    yield sim_panda
+
+
+@patch("ophyd_async.core.device_save_loader.save_to_yaml")
+async def test_save_panda(mock_save_to_yaml, sim_panda, RE: RunEngine):
+    RE(save_device(sim_panda, "path", sorter=phase_sorter))
+    mock_save_to_yaml.assert_called_once()
+    assert mock_save_to_yaml.call_args[0] == (
+        [
+            {"phase_1_signal_units": 0},
+            {
+                "pulse.1.delay": 0.0,
+                "pulse.1.width": 0.0,
+                "seq.1.table": {},
+                "seq.1.active": False,
+            },
+        ],
+        "path",
+    )