Make run_number a metadata key in period-per-point mode

Author: Tom-Willemsen

src/ibex_bluesky_core/plans/__init__.py

@@ -6,6 +6,7 @@
 import bluesky.plans as bp
 import matplotlib.pyplot as plt
 from bluesky import plan_stubs as bps
+from bluesky.preprocessors import inject_md_decorator
 from bluesky.protocols import NamedMovable
 from bluesky.utils import Msg
 from matplotlib.axes import Axes
@@ -31,6 +32,17 @@
 ]
 
 
+def _get_additional_md(
+    dae: "SimpleDae", *, periods: bool, save_run: bool
+) -> Generator[Msg, None, dict[str, Any]]:
+    if periods and save_run:
+        run_number = yield from bps.rd(dae.current_or_next_run_number_str)
+        return {"run_number": run_number}
+    else:
+        yield from bps.null()
+        return {}
+
+
 def scan(  # noqa: PLR0913
     dae: "SimpleDae",
     block: NamedMovable[float],
@@ -66,7 +78,10 @@ def scan(  # noqa: PLR0913
 
     icc = _set_up_fields_and_icc(block, dae, model, periods, save_run, ax)
 
+    additional_md = yield from _get_additional_md(dae, periods=periods, save_run=save_run)
+
     @icc
+    @inject_md_decorator(additional_md)
     def _inner() -> Generator[Msg, None, None]:
         if rel:
             plan = bp.rel_scan
@@ -148,7 +163,10 @@ def adaptive_scan(  # noqa: PLR0913, PLR0917
 
     icc = _set_up_fields_and_icc(block, dae, model, periods, save_run, ax)
 
+    additional_md = yield from _get_additional_md(dae, periods=periods, save_run=save_run)
+
     @icc
+    @inject_md_decorator(additional_md)
     def _inner() -> Generator[Msg, None, None]:
         if rel:
             plan = bp.rel_adaptive_scan

tests/plans/test_init.py

@@ -1,6 +1,9 @@
 # pyright: reportMissingParameterType=false
+import functools
+from typing import Any
 from unittest.mock import patch
 
+import bluesky.utils
 import pytest
 from bluesky.preprocessors import run_decorator
 from ophyd_async.plan_stubs import ensure_connected
@@ -121,6 +124,36 @@ def test_scan_does_normal_scan_when_relative_false(RE, dae, block):
     assert count == bp_scan.call_args[1]["num"]
 
 
+@pytest.mark.parametrize(
+    "scan_func",
+    [
+        functools.partial(scan, start=1, stop=2, num=2),
+        functools.partial(adaptive_scan, start=1, stop=2, min_step=1, max_step=2, target_delta=1),
+    ],
+)
+def test_if_in_periods_mode_and_run_saved_then_scan_start_doc_contains_run_number(
+    RE, dae, block, scan_func
+):
+    set_mock_value(dae.current_or_next_run_number_str, "12345678")
+
+    start_doc: dict[str, Any] | None = None
+
+    def _cb(typ, doc):
+        if typ == "start":
+            nonlocal start_doc
+            start_doc = doc
+
+    with (
+        patch("ibex_bluesky_core.plans.ensure_connected"),
+    ):
+        # Scan fails because DAE isn't set up right... but it still emits a start doc so that's fine
+        with pytest.raises(bluesky.utils.FailedStatus):
+            RE(scan_func(dae, block, rel=False, periods=True, save_run=True), _cb)
+
+    assert start_doc is not None
+    assert start_doc.get("run_number") == "12345678"
+
+
 def test_scan_does_relative_scan_when_relative_true(RE, dae, block):
     start = 0
     stop = 2