Move plans and stubs from BlueAPI into repository structure

- Enforce mypy on untyped functions
- Add documentation to packages explaining the functionality and requirements of plans, stubs
- Use __export__ to explicitly export "standalone" stubs
- Add tests to check compliance of plans, stubs to requirements
- Ensure loop for 3.10 Future

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/blueapi/plans/plans.py → src/dls_bluesky_core/plans/scanspec.py

@@ -1,13 +1,19 @@
 import operator
 from functools import reduce
-from typing import Any, List, Mapping, Optional, Union
+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 blueapi.core import MsgGenerator
+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(
@@ -71,43 +77,3 @@ def _scanspec_to_cycler(spec: Spec[str], axes: Mapping[str, Movable]) -> Cycler:
     # 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()))
-
-
-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/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/blueapi/plans/stubs.py → src/dls_bluesky_core/stubs/wrapped.py

@@ -1,13 +1,16 @@
 import itertools
-from typing import Any, Mapping, Optional, TypeVar
+from typing import Annotated, Any, Mapping, Optional, TypeVar
 
 import bluesky.plan_stubs as bps
 from bluesky.protocols import Movable
 
-from blueapi.core import MsgGenerator
+from dls_bluesky_core.core import MsgGenerator
 
-Group = str
+"""
+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")
 
 

tests/plans/test_compliance.py

@@ -0,0 +1,68 @@
+import inspect
+from types import ModuleType
+from typing import Any, List, Mapping, Optional, get_type_hints
+
+import dls_bluesky_core.plans as plans
+import dls_bluesky_core.stubs as stubs
+from dls_bluesky_core.core import MsgGenerator, PlanGenerator
+
+
+def is_bluesky_plan_generator(func: Any) -> bool:
+    try:
+        return get_type_hints(func).get("return") is MsgGenerator
+    except TypeError:
+        # get_type_hints fails on some objects (such as Union or Optional)
+        return False
+
+
+def get_all_available_generators(mod: ModuleType):
+    def get_named_subset(names: List[str]):
+        for name in names:
+            yield getattr(mod, name)
+
+    if "__export__" in mod.__dict__:
+        yield from get_named_subset(getattr(mod, "__export__"))
+    elif "__all__" in mod.__dict__:
+        yield from get_named_subset(getattr(mod, "__all__"))
+    else:
+        for name, value in mod.__dict__.items():
+            if not name.startswith("_"):
+                yield value
+
+
+def assert_hard_requirements(plan: PlanGenerator, signature: inspect.Signature):
+    assert plan.__doc__ is not None, f"'{plan.__name__}' has no docstring"
+    for parameter in signature.parameters.values():
+        assert (
+            parameter.kind is not parameter.VAR_POSITIONAL
+            and parameter.kind is not parameter.VAR_KEYWORD
+        ), f"'{plan.__name__}' has variadic arguments"
+
+
+def assert_metadata_requirements(plan: PlanGenerator, signature: inspect.Signature):
+    assert (
+        "metadata" in signature.parameters
+    ), f"'{plan.__name__}' does not allow metadata"
+    metadata = signature.parameters["metadata"]
+    assert (
+        metadata.annotation is Optional[Mapping[str, Any]]
+        and metadata.default is not inspect.Parameter.empty
+    ), f"'{plan.__name__}' metadata is not optional"
+    assert metadata.default is None, f"'{plan.__name__}' metadata default is mutable"
+
+
+def test_plans_comply():
+    for plan in get_all_available_generators(plans):
+        if is_bluesky_plan_generator(plan):
+            signature = inspect.Signature.from_callable(plan)
+            assert_hard_requirements(plan, signature)
+            assert_metadata_requirements(plan, signature)
+
+
+def test_stubs_comply():
+    for plan in get_all_available_generators(stubs):
+        if is_bluesky_plan_generator(plan):
+            signature = inspect.Signature.from_callable(plan)
+            assert_hard_requirements(plan, signature)
+            if "metadata" in signature.parameters:
+                assert_metadata_requirements(plan, signature)

tests/plans/test_scanspec_metadata.py

@@ -1,11 +1,18 @@
-from asyncio import Future
+from asyncio import Future, new_event_loop
 
 import pytest
 from bluesky import RunEngine
 from ophyd.sim import Syn2DGauss, SynAxis, SynGauss
 from scanspec.specs import Line, Spiral
 
-from blueapi.plans import scan
+from dls_bluesky_core.plans import scan
+
+
+@pytest.fixture(scope="session")
+def loop():
+    loop = new_event_loop()
+    yield loop
+    loop.stop()
 
 
 @pytest.fixture
@@ -30,7 +37,7 @@ def set_result(_, doc):
     return run_engine.subscribe(name=name, func=set_result)
 
 
-def test_metadata_of_simple_spec(run_engine, x):
+def test_metadata_of_simple_spec(run_engine, x, loop):
     det = SynGauss(
         name="det",
         motor=x,
@@ -41,7 +48,7 @@ def test_metadata_of_simple_spec(run_engine, x):
     )
     spec = Line(axis=x.name, start=1, stop=2, num=3)
 
-    start_future = Future()
+    start_future = loop.create_future()
     tok = capture_document_return_token(start_future, run_engine, "start")
     run_engine(scan([det], {x.name: x}, spec))
     run_engine.unsubscribe(tok)
@@ -60,7 +67,7 @@ def test_metadata_of_simple_spec(run_engine, x):
     assert start_document["detectors"] == [det.name]
 
 
-def test_metadata_of_spiral_spec(run_engine, x, y):
+def test_metadata_of_spiral_spec(run_engine, x, y, loop):
     det = Syn2DGauss(
         name="det",
         motor0=x,
@@ -73,7 +80,7 @@ def test_metadata_of_spiral_spec(run_engine, x, y):
     )
     spec = Spiral.spaced(x.name, y.name, 0, 0, 5, 1)
 
-    start_future = Future()
+    start_future = loop.create_future()
     tok = capture_document_return_token(start_future, run_engine, "start")
     run_engine(scan([det], {x.name: x, y.name: y}, spec))
     run_engine.unsubscribe(tok)