dls-dodal 1.35.0__py3-none-any.whl → 1.36.1a0__py3-none-any.whl

dodal/plan_stubs/wrapped.py

@@ -0,0 +1,150 @@
+import itertools
+from collections.abc import Mapping
+from typing import Annotated, Any
+
+import bluesky.plan_stubs as bps
+from bluesky.protocols import Movable
+from bluesky.utils import MsgGenerator
+
+"""
+Wrappers for Bluesky built-in plan stubs with type hinting
+"""
+
+Group = Annotated[str, "String identifier used by 'wait' or stubs that await"]
+
+
+# After bluesky 1.14, bounds for stubs that move can be narrowed
+# https://github.com/bluesky/bluesky/issues/1821
+def set_absolute(
+    movable: Movable, value: Any, group: Group | None = 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 (Group | None, 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: Any, group: Group | None = 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 (Group | None, 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: Group | None = None) -> MsgGenerator:
+    """
+    Move a device, wrapper for `bp.mv`.
+
+    Args:
+        moves (Mapping[Movable, T]): Mapping of Movables to target positions
+        group (Group | None, optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (
+        # type ignore until https://github.com/bluesky/bluesky/issues/1809
+        yield from bps.mv(*itertools.chain.from_iterable(moves.items()), group=group)  # type: ignore
+    )
+
+
+def move_relative(
+    moves: Mapping[Movable, Any], group: Group | None = None
+) -> MsgGenerator:
+    """
+    Move a device relative to its current position, wrapper for `bp.mvr`.
+
+    Args:
+        moves (Mapping[Movable, T]): Mapping of Movables to target deltas
+        group (Group | None, optional): The message group to associate with the
+                                           setting, for sequencing. Defaults to None.
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (
+        # type ignore until https://github.com/bluesky/bluesky/issues/1809
+        yield from bps.mvr(*itertools.chain.from_iterable(moves.items()), group=group)  # type: ignore
+    )
+
+
+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: Group | None = None,
+    timeout: float | None = None,
+) -> MsgGenerator:
+    """
+    Wait for a group status to complete, wrapper for `bp.wait`.
+    Does not expose move_on, as when used as a stub will not fail on Timeout.
+
+    Args:
+        group (Group | None, optional): The name of the group to wait for, defaults
+                                           to None, in which case waits for all
+                                           groups that have not yet been awaited.
+        timeout (float | None, default=None): a timeout in seconds
+
+
+    Returns:
+        MsgGenerator: Plan
+
+    Yields:
+        Iterator[MsgGenerator]: Bluesky messages
+    """
+
+    return (yield from bps.wait(group, timeout=timeout))

dodal/plans/__init__.py

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

dodal/plans/scanspec.py

@@ -0,0 +1,66 @@
+import operator
+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.data_session 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],
+        Field(description="ScanSpec modelling the path of the scan"),
+    ],
+    metadata: dict[str, Any] | None = None,
+) -> MsgGenerator:
+    """Generic plan for reading `detectors` at every point of a ScanSpec `Spec`.
+    A `Spec` is an N-dimensional path.
+    """
+    # TODO: https://github.com/bluesky/scanspec/issues/154
+    # support Static.duration: Spec[Literal["DURATION"]]
+
+    _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(tuple(detectors), _as_cycler(spec), md=_md)
+
+
+def _as_cycler(spec: Spec[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
+
+    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()))

dodal/plans/wrapped.py

@@ -0,0 +1,57 @@
+from collections.abc import Sequence
+from typing import Annotated, Any
+
+import bluesky.plans as bp
+from bluesky.protocols import Readable
+from pydantic import Field, NonNegativeFloat, validate_call
+
+from dodal.common import MsgGenerator
+from dodal.plan_stubs.data_session import attach_data_session_metadata_decorator
+
+"""This module wraps plan(s) from bluesky.plans until required handling for them is
+moved into bluesky or better handled in downstream services.
+
+Required decorators are installed on plan import
+https://github.com/DiamondLightSource/blueapi/issues/474
+
+Non-serialisable fields are ignored when they are optional
+https://github.com/DiamondLightSource/blueapi/issues/711
+
+We may also need other adjustments for UI purposes, e.g.
+Forcing uniqueness or orderedness of Readables
+Limits and metadata (e.g. units)
+"""
+
+
+@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[
+        NonNegativeFloat | Sequence[NonNegativeFloat],
+        Field(
+            description="Delay between readings: if tuple, 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"},
+        ),
+    ] = 0.0,
+    metadata: dict[str, Any] | None = None,
+) -> MsgGenerator:
+    """Reads from a number of devices.
+    Wraps bluesky.plans.count(det, num, delay, md=metadata) exposing only serializable
+    parameters and metadata."""
+    if isinstance(delay, Sequence):
+        assert (
+            len(delay) == num - 1
+        ), f"Number of delays given must be {num - 1}: was given {len(delay)}"
+    metadata = metadata or {}
+    metadata["shape"] = (num,)
+    yield from bp.count(tuple(detectors), num, delay=delay, md=metadata)

dodal/utils.py

@@ -1,3 +1,4 @@
+import functools
 import importlib
 import inspect
 import os
@@ -6,13 +7,14 @@ import socket
 import string
 from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass
-from functools import wraps
+from functools import update_wrapper, wraps
 from importlib import import_module
 from inspect import signature
 from os import environ
 from types import ModuleType
 from typing import (
     Any,
+    Generic,
     Protocol,
     TypeGuard,
     TypeVar,
@@ -35,6 +37,7 @@ from bluesky.protocols import (
     Triggerable,
     WritesExternalAssets,
 )
+from bluesky.run_engine import call_in_bluesky_event_loop
 from ophyd.device import Device as OphydV1Device
 from ophyd_async.core import Device as OphydV2Device
 
@@ -99,6 +102,8 @@ class BeamlinePrefix:
 
 
 T = TypeVar("T", bound=AnyDevice)
+D = TypeVar("D", bound=OphydV2Device)
+SkipType = bool | Callable[[], bool]
 
 
 def skip_device(precondition=lambda: True):
@@ -114,6 +119,91 @@ def skip_device(precondition=lambda: True):
     return decorator
 
 
+class DeviceInitializationController(Generic[D]):
+    def __init__(
+        self,
+        factory: Callable[[], D],
+        use_factory_name: bool,
+        timeout: float,
+        mock: bool,
+        skip: SkipType,
+    ):
+        self._factory: Callable[[], D] = functools.cache(factory)
+        self._use_factory_name = use_factory_name
+        self._timeout = timeout
+        self._mock = mock
+        self._skip = skip
+        update_wrapper(self, factory)
+
+    @property
+    def skip(self) -> bool:
+        return self._skip() if callable(self._skip) else self._skip
+
+    def cache_clear(self) -> None:
+        """Clears the controller's internal cached instance of the device, if present.
+        Noop if not."""
+
+        # Functools adds the cache_clear function via setattr so the type checker
+        # does not pick it up.
+        self._factory.cache_clear()  # type: ignore
+
+    def __call__(
+        self,
+        connect_immediately: bool = False,
+        name: str | None = None,
+        connection_timeout: float | None = None,
+        mock: bool | None = None,
+    ) -> D:
+        """Returns an instance of the Device the wrapped factory produces: the same
+        instance will be returned if this method is called multiple times, and arguments
+        may be passed to override this Controller's configuration.
+        Once the device is connected, the value of mock must be consistent, or connect
+        must be False.
+
+
+        Args:
+            connect_immediately (bool, default False): whether to call connect on the
+              device before returning it- connect is idempotent for ophyd-async devices.
+              Not connecting to the device allows for the instance to be created prior
+              to the RunEngine event loop being configured or for connect to be called
+              lazily e.g. by the `ensure_connected` stub.
+            name (str | None, optional): an override name to give the device, which is
+              also used to name its children. Defaults to None, which does not name the
+              device unless the device has no name and this Controller is configured to
+              use_factory_name, which propagates the name of the wrapped factory
+              function to the device instance.
+            connection_timeout (float | None, optional): an override timeout length in
+              seconds for the connect method, if it is called. Defaults to None, which
+              defers to the timeout configured for this Controller: the default uses
+              ophyd_async's DEFAULT_TIMEOUT.
+            mock (bool | None, optional): overrides whether to connect to Mock signal
+              backends, if connect is called. Defaults to None, which uses the mock
+              parameter of this Controller. This value must be used consistently when
+              connect is called on the Device.
+
+        Returns:
+            D: a singleton instance of the Device class returned by the wrapped factory.
+        """
+        device = self._factory()
+
+        if connect_immediately:
+            call_in_bluesky_event_loop(
+                device.connect(
+                    timeout=connection_timeout
+                    if connection_timeout is not None
+                    else self._timeout,
+                    mock=mock if mock is not None else self._mock,
+                )
+            )
+
+        if name:
+            device.set_name(name)
+        elif not device.name and self._use_factory_name:
+            device.set_name(self._factory.__name__)
+
+        return device
+
+
 def make_device(
     module: str | ModuleType,
     device_name: str,
@@ -206,7 +296,33 @@ def invoke_factories(
         dependent_name = leaves.pop()
         params = {name: devices[name] for name in dependencies[dependent_name]}
         try:
-            devices[dependent_name] = factories[dependent_name](**params, **kwargs)
+            factory = factories[dependent_name]
+            if isinstance(factory, DeviceInitializationController):
+                # For now we translate the old-style parameters that
+                # device_instantiation expects. Once device_instantiation is gone and
+                # replaced with DeviceInitializationController we can formalise the
+                # API of make_all_devices and make these parameters explicit.
+                # https://github.com/DiamondLightSource/dodal/issues/844
+                mock = kwargs.get(
+                    "mock",
+                    kwargs.get(
+                        "fake_with_ophyd_sim",
+                        False,
+                    ),
+                )
+                connect_immediately = kwargs.get(
+                    "connect_immediately",
+                    kwargs.get(
+                        "wait_for_connection",
+                        False,
+                    ),
+                )
+                devices[dependent_name] = factory(
+                    mock=mock,
+                    connect_immediately=connect_immediately,
+                )
+            else:
+                devices[dependent_name] = factory(**params, **kwargs)
         except Exception as e:
             exceptions[dependent_name] = e
 
@@ -268,6 +384,8 @@ def collect_factories(
 
 
 def _is_device_skipped(func: AnyDeviceFactory) -> bool:
+    if isinstance(func, DeviceInitializationController):
+        return func.skip
     return getattr(func, "__skip__", False)
 
 
@@ -301,6 +419,37 @@ def is_v1_device_type(obj: type[Any]) -> bool:
     return is_class and follows_protocols and not is_v2_device_type(obj)
 
 
+def filter_ophyd_devices(
+    devices: Mapping[str, AnyDevice],
+) -> tuple[Mapping[str, OphydV1Device], Mapping[str, OphydV2Device]]:
+    """
+    Split a dictionary of ophyd and ophyd-async devices
+    (i.e. the output of make_all_devices) into 2 separate dictionaries of the
+    different types. Useful when special handling is needed for each type of device.
+
+    Args:
+        devices: Dictionary of device name to ophyd or ophyd-async device.
+
+    Raises:
+        ValueError: If anything in the dictionary doesn't come from either library.
+
+    Returns:
+        Tuple of two dictionaries, one mapping names to ophyd devices and one mapping
+        names to ophyd-async devices.
+    """
+
+    ophyd_devices = {}
+    ophyd_async_devices = {}
+    for name, device in devices.items():
+        if isinstance(device, OphydV1Device):
+            ophyd_devices[name] = device
+        elif isinstance(device, OphydV2Device):
+            ophyd_async_devices[name] = device
+        else:
+            raise ValueError(f"{name}: {device} is not an ophyd or ophyd-async device")
+    return ophyd_devices, ophyd_async_devices
+
+
 def get_beamline_based_on_environment_variable() -> ModuleType:
     """
     Gets the dodal module for the current beamline, as specified by the