Move plans and stubs from BlueAPI repository

pyproject.toml

@@ -13,6 +13,9 @@ classifiers = [
 ]
 description = "Common Diamond specific Bluesky plans and functions"
 dependencies = [
+    "blueapi",
+    "ophyd",
+    "scanspec"
 ] # Add project dependencies here, e.g. ["click", "numpy"]
 dynamic = ["version"]
 license.file = "LICENSE"
@@ -53,6 +56,10 @@ write_to = "src/dls_bluesky_core/_version.py"
 
 [tool.mypy]
 ignore_missing_imports = true # Ignore missing stubs in imported modules
+[[tool.mypy.overrides]]
+# Enforce disallow_untyped_degs on all src/ but not tests/
+module = 'dls_bluesky_core.*'
+disallow_untyped_defs = true
 
 [tool.isort]
 float_to_top = true

src/dls_bluesky_core/__main__.py

@@ -1,14 +1,15 @@
 from argparse import ArgumentParser
+from typing import Optional, Sequence
 
 from . import __version__
 
 __all__ = ["main"]
 
 
-def main(args=None):
+def main(args: Optional[Sequence[str]] = None) -> None:
     parser = ArgumentParser()
     parser.add_argument("-v", "--version", action="version", version=__version__)
-    args = parser.parse_args(args)
+    parser.parse_args(args)
 
 
 # test with: python -m dls_bluesky_core

src/dls_bluesky_core/core/__init__.py

@@ -0,0 +1,6 @@
+from .types import MsgGenerator, PlanGenerator
+
+__all__ = [
+    "MsgGenerator",
+    "PlanGenerator",
+]

src/dls_bluesky_core/core/types.py

@@ -0,0 +1,8 @@
+from typing import Any, Callable, Generator
+
+from bluesky import Msg
+
+#  'A true "plan", usually the output of a generator function'
+MsgGenerator = Generator[Msg, Any, None]
+#  'A function that generates a plan'
+PlanGenerator = Callable[..., MsgGenerator]

src/dls_bluesky_core/plans/__init__.py

@@ -0,0 +1,23 @@
+from .scanspec import scan
+from .wrapped import count
+
+"""
+This package is intended to hold MsgGenerator functions which act as self-contained
+experiments: they start runs, collect data, and close the runs. While they may be used
+as building blocks for larger nested plans, they are primarily intended to be run as-is,
+and any common functionality which may be useful to multiple plans extracted to stubs/.
+
+Plans:
+- Must have type hinted arguments, Should use the loosest sensible bounds
+- Must have docstrings describing behaviour and arguments of the function
+- Must not have variadic args or kwargs, Should pass collections instead
+- Must have optional argument named 'metadata' to add metadata to run(s)
+- Must add 'plan_args' to metadata with complete representation including defaults, None
+- Must add 'detectors', 'motors' metadata with list of names of relevant devices
+- Should pass 'shape' to metadata if the run's shape is knowable
+"""
+
+__all__ = [
+    "count",
+    "scan",
+]

src/dls_bluesky_core/plans/scanspec.py

@@ -0,0 +1,79 @@
+import operator
+from functools import reduce
+from typing import Any, List, Mapping, Optional
+
+import bluesky.plans as bp
+from bluesky.protocols import Movable, Readable
+from cycler import Cycler, cycler
+from scanspec.specs import Spec
+
+from dls_bluesky_core.core import MsgGenerator
+
+"""
+Plans related to the use of the `ScanSpec https://github.com/dls-controls/scanspec`
+library for constructing arbitrarily complex N-dimensional trajectories, similar to
+Diamond's "mapping scans" using ScanPointGenerator.
+"""
+
+
+def scan(
+    detectors: List[Readable],
+    axes_to_move: Mapping[str, Movable],
+    spec: Spec[str],
+    metadata: Optional[Mapping[str, Any]] = None,
+) -> MsgGenerator:
+    """
+    Scan wrapping `bp.scan_nd`
+
+    Args:
+        detectors: List of readable devices, will take a reading at
+                                    each point
+        axes_to_move: All axes involved in this scan, names and
+            objects
+        spec: ScanSpec modelling the path of the scan
+        metadata: Key-value metadata to include
+                                                          in exported data, defaults to
+                                                          None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    _md = {
+        "plan_args": {
+            "detectors": list(map(repr, detectors)),
+            "axes_to_move": {k: repr(v) for k, v in axes_to_move.items()},
+            "spec": repr(spec),
+        },
+        "plan_name": "scan",
+        "shape": spec.shape(),
+        **(metadata or {}),
+    }
+
+    cycler = _scanspec_to_cycler(spec, axes_to_move)
+    yield from bp.scan_nd(detectors, cycler, md=_md)
+
+
+def _scanspec_to_cycler(spec: Spec[str], axes: Mapping[str, Movable]) -> Cycler:
+    """
+    Convert a scanspec to a cycler for compatibility with legacy Bluesky plans such as
+    `bp.scan_nd`. Use the midpoints of the scanspec since cyclers are normally used
+    for software triggered scans.
+
+    Args:
+        spec: A scanspec
+        axes: Names and axes to move
+
+    Returns:
+        Cycler: A new cycler
+    """
+
+    midpoints = spec.frames().midpoints
+    midpoints = {axes[name]: points for name, points in midpoints.items()}
+
+    # Need to "add" the cyclers for all the axes together. The code below is
+    # effectively: cycler(motor1, [...]) + cycler(motor2, [...]) + ...
+    return reduce(operator.add, map(lambda args: cycler(*args), midpoints.items()))

src/dls_bluesky_core/plans/wrapped.py

@@ -0,0 +1,50 @@
+from typing import Any, List, Mapping, Optional, Union
+
+import bluesky.plans as bp
+from bluesky.protocols import Readable
+
+from dls_bluesky_core.core import MsgGenerator
+
+"""
+Wrappers for Bluesky built-in plans with type hinting and renamed metadata
+"""
+
+
+def count(
+    detectors: List[Readable],
+    num: int = 1,
+    delay: Optional[Union[float, List[float]]] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+) -> MsgGenerator:
+    """
+    Take `n` readings from a device
+
+    Args:
+        detectors (List[Readable]): Readable devices to read
+        num (int, optional): Number of readings to take. Defaults to 1.
+        delay (Optional[Union[float, List[float]]], optional): Delay between readings.
+                                                               Defaults to None.
+        metadata (Optional[Mapping[str, Any]], optional): Key-value metadata to include
+                                                          in exported data.
+                                                          Defaults to None.
+
+    Returns:
+        MsgGenerator: _description_
+
+    Yields:
+        Iterator[MsgGenerator]: _description_
+    """
+    plan_args = (
+        {  # If bp.count added delay to plan_args, we could remove all md handling
+            "detectors": list(map(repr, detectors)),
+            "num": num,
+            "delay": delay,
+        }
+    )
+
+    _md = {
+        "plan_args": plan_args,
+        **(metadata or {}),
+    }
+
+    yield from bp.count(detectors, num, delay=delay, md=_md)

src/dls_bluesky_core/stubs/__init__.py

@@ -0,0 +1,41 @@
+from typing import List
+
+from .wrapped import move, move_relative, set_absolute, set_relative, sleep, wait
+
+"""
+This package is intended to hold MsgGenerator functions which are not self-contained
+data collections: while they may start runs, collect data, or close runs, they are
+blocks for larger nested plans, and may not make sense to be run as-is. Functions that
+may make sense as isolated blocks of functionality (e.g. moving a motor) should be added
+to the __export__ list: without this list, it is assumed that all MsgGenerator functions
+in the package should be imported by any services which respect it.
+
+Functions that yield from multiple stubs and offer a complete workflow
+should be moved to plans/.
+
+This package should never have a dependency on plans/.
+
+Stubs:
+- Must have type hinted arguments, Should use the loosest sensible bounds
+- Must have docstrings describing behaviour and arguments of the function
+- Must not have variadic args or kwargs, Should pass collections instead
+- Allow metadata to be propagated through if calling other stubs that take metadata
+"""
+
+__export__: List[str] = [  # Available for import to services
+    "set_absolute",
+    "set_relative",
+    "move",
+    "move_relative",
+    "sleep",
+    "wait",
+]
+
+__all__: List[str] = [  # Available for import by other modules
+    "set_absolute",
+    "set_relative",
+    "move",
+    "move_relative",
+    "sleep",
+    "wait",
+]

src/dls_bluesky_core/stubs/wrapped.py

@@ -0,0 +1,141 @@
+import itertools
+from typing import Annotated, Any, Mapping, Optional, TypeVar
+
+import bluesky.plan_stubs as bps
+from bluesky.protocols import Movable
+
+from dls_bluesky_core.core import MsgGenerator
+
+"""
+Wrappers for Bluesky built-in plan stubs with type hinting
+"""
+
+Group = Annotated[str, "String identifier used by 'wait' or stubs that await"]
+T = TypeVar("T")
+
+
+def set_absolute(
+    movable: Movable, value: T, group: Optional[Group] = None, wait: bool = False
+) -> MsgGenerator:
+    """
+    Set a device, wrapper for `bp.abs_set`.
+
+    Args:
+        movable (Movable): The device to set
+        value (T): The new value
+        group (Optional[Group], optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+        wait (bool, optional): The group should wait until all setting is complete
+                               (e.g. a motor has finished moving). Defaults to False.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (yield from bps.abs_set(movable, value, group=group, wait=wait))
+
+
+def set_relative(
+    movable: Movable, value: T, group: Optional[Group] = None, wait: bool = False
+) -> MsgGenerator:
+    """
+    Change a device, wrapper for `bp.rel_set`.
+
+    Args:
+        movable (Movable): The device to set
+        value (T): The new value
+        group (Optional[Group], optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+        wait (bool, optional): The group should wait until all setting is complete
+                               (e.g. a motor has finished moving). Defaults to False.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (yield from bps.rel_set(movable, value, group=group, wait=wait))
+
+
+def move(moves: Mapping[Movable, Any], group: Optional[Group] = None) -> MsgGenerator:
+    """
+    Move a device, wrapper for `bp.mv`.
+
+    Args:
+        moves (Mapping[Movable, Any]): Mapping of Movables to target positions
+        group (Optional[Group], optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (
+        yield from bps.mv(*itertools.chain.from_iterable(moves.items()), group=group)
+    )
+
+
+def move_relative(
+    moves: Mapping[Movable, Any], group: Optional[Group] = None
+) -> MsgGenerator:
+    """
+    Move a device relative to its current position, wrapper for `bp.mvr`.
+
+    Args:
+        moves (Mapping[Movable, Any]): Mapping of Movables to target deltas
+        group (Optional[Group], optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (
+        yield from bps.mvr(*itertools.chain.from_iterable(moves.items()), group=group)
+    )
+
+
+def sleep(time: float) -> MsgGenerator:
+    """
+    Suspend all action for a given time, wrapper for `bp.sleep`
+
+    Args:
+        time (float): Time to wait in seconds
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (yield from bps.sleep(time))
+
+
+def wait(group: Optional[Group] = None) -> MsgGenerator:
+    """
+    Wait for a group status to complete, wrapper for `bp.wait`
+
+    Args:
+        group (Optional[Group], optional): The name of the group to wait for, defaults
+                                           to None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (yield from bps.wait(group))