Add wrapped count plan and generic ScanSpec scan from dls-bluesky-core

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "aiofiles",
     "aiohttp",
     "redis",
+    "scanspec",
 ]
 
 dynamic = ["version"]

src/dodal/plans/__init__.py

@@ -0,0 +1,4 @@
+from .scanspec import spec_scan
+from .wrapped import count
+
+__all__ = ["count", "spec_scan"]

src/dodal/plans/scanspec.py

@@ -0,0 +1,59 @@
+import operator
+from collections.abc import Mapping
+from functools import reduce
+from typing import Annotated, Any
+
+import bluesky.plans as bp
+from bluesky.protocols import Movable, Readable
+from cycler import Cycler, cycler
+from pydantic import Field, validate_call
+from scanspec.specs import Spec
+
+from dodal.common import MsgGenerator
+from dodal.plan_stubs import attach_data_session_metadata_decorator
+
+
+@attach_data_session_metadata_decorator()
+@validate_call(config={"arbitrary_types_allowed": True})
+def spec_scan(
+    detectors: Annotated[
+        set[Readable],
+        Field(
+            description="Set of readable devices, will take a reading at each point, \
+            in addition to any Movables in the Spec",
+        ),
+    ],
+    spec: Annotated[Spec[Movable], "ScanSpec modelling the path of the scan"],
+    metadata: Mapping[str, Any] | None = None,
+) -> MsgGenerator:
+    """Generic plan for reading `detectors` at every point of a ScanSpec `spec`."""
+    _md = {
+        "plan_args": {
+            "detectors": {det.name for det in detectors},
+            "spec": repr(spec),
+        },
+        "plan_name": "spec_scan",
+        "shape": spec.shape(),
+        **(metadata or {}),
+    }
+
+    yield from bp.scan_nd(detectors, _as_cycler(spec), md=_md)
+
+
+def _as_cycler(spec: Spec[str]) -> 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
+
+    Returns:
+        Cycler: A new cycler
+    """
+
+    midpoints = spec.frames().midpoints
+    # Need to "add" the cyclers for all the axes together. The code below is
+    # effectively: cycler(motor1, [...]) + cycler(motor2, [...]) + ...
+    return reduce(operator.add, (cycler(*args) for args in midpoints.items()))

src/dodal/plans/wrapped.py

@@ -0,0 +1,38 @@
+from collections.abc import Mapping
+from typing import Annotated, Any
+
+import bluesky.plans as bp
+from bluesky.protocols import Readable
+from pydantic import Field, PositiveFloat, validate_call
+
+from dodal.common import MsgGenerator
+from dodal.plan_stubs import attach_data_session_metadata_decorator
+
+
+@attach_data_session_metadata_decorator()
+@validate_call(config={"arbitrary_types_allowed": True})
+def count(
+    detectors: Annotated[
+        set[Readable],
+        Field(
+            description="Set of readable devices, will take a reading at each point",
+            min_length=1,
+        ),
+    ],
+    num: Annotated[int, Field(description="Number of frames to collect", ge=1)] = 1,
+    delay: Annotated[
+        PositiveFloat | list[PositiveFloat] | None,
+        Field(
+            description="Delay between readings: if list, len(delay) == num - 1 and \
+            the delays are between each point, if value or None is the delay for every \
+            gap",
+            json_schema_extra={"units": "s"},
+        ),
+    ] = None,
+    metadata: Mapping[str, Any] | None = None,
+) -> MsgGenerator:
+    if isinstance(delay, list):
+        assert (
+            delays := len(delay)
+        ) == num - 1, f"Number of delays given must be {num - 1}: was given {delays} "
+    yield from bp.count(detectors, num, delay=delay, md=metadata or {})