Slits: report blade progress with WatchableAsyncStatus

src/hextools/photon_delivery_system/slits.py

@@ -5,9 +5,10 @@
 
 from ophyd_async.core import (
     AsyncMovable,
-    AsyncStatus,
     StandardReadable,
+    WatchableAsyncStatus,
     derived_signal_rw,
+    forward_watcher_updates,
 )
 from ophyd_async.core import StandardReadableFormat as Format
 from ophyd_async.epics.core import EpicsDevice
@@ -106,10 +107,10 @@ def _get_center(self, low: float, high: float) -> float:
         """
         return (low + high) / 2
 
-    @AsyncStatus.wrap
+    @WatchableAsyncStatus.wrap
     async def _set_gap(
         self, low_motor: AsyncEpicsMotor, high_motor: AsyncEpicsMotor, gap: float
-    ) -> None:
+    ):
         """Set the gap while holding the center fixed.
 
         Parameters
@@ -126,15 +127,17 @@ async def _set_gap(
             high_motor.user_readback.get_value(),
         )
         center = (low + high) / 2
-        await asyncio.gather(
+        statuses = [
             low_motor.set(center - gap / 2),
             high_motor.set(center + gap / 2),
-        )
+        ]
+        async for update in forward_watcher_updates(statuses, self.name, combine=False):
+            yield update
 
-    @AsyncStatus.wrap
+    @WatchableAsyncStatus.wrap
     async def _set_center(
         self, low_motor: AsyncEpicsMotor, high_motor: AsyncEpicsMotor, center: float
-    ) -> None:
+    ):
         """Set the center while holding the gap fixed.
 
         Parameters
@@ -151,12 +154,14 @@ async def _set_center(
             high_motor.user_readback.get_value(),
         )
         shift = center - (low + high) / 2
-        await asyncio.gather(
+        statuses = [
             low_motor.set(low + shift),
             high_motor.set(high + shift),
-        )
+        ]
+        async for update in forward_watcher_updates(statuses, self.name, combine=False):
+            yield update
 
-    @AsyncStatus.wrap
+    @WatchableAsyncStatus.wrap
     async def set(
         self,
         value: tuple[tuple[float, float], tuple[float, float]]
@@ -180,12 +185,16 @@ async def set(
             )
 
         # Set gaps first to preserve centers, then set centers to preserve gaps.
-        await asyncio.gather(
-            self.horizontal_gap.set(h_gap),
-            self.vertical_gap.set(v_gap),
-        )
-
-        await asyncio.gather(
-            self.horizontal_center.set(h_center),
-            self.vertical_center.set(v_center),
-        )
+        gaps = [
+            self._set_gap(self.inboard, self.outboard, h_gap),
+            self._set_gap(self.bottom, self.top, v_gap),
+        ]
+        async for update in forward_watcher_updates(gaps, self.name, combine=False):
+            yield update
+
+        centers = [
+            self._set_center(self.inboard, self.outboard, h_center),
+            self._set_center(self.bottom, self.top, v_center),
+        ]
+        async for update in forward_watcher_updates(centers, self.name, combine=False):
+            yield update

tests/photon_delivery_system/test_slits.py

@@ -96,3 +96,20 @@ async def test_slits_set_all_gaps_and_centers_via_mv(RE: RunEngine):
     assert await slits.horizontal_center.get_value() == pytest.approx(10.0)
     assert await slits.vertical_gap.get_value() == pytest.approx(4.0)
     assert await slits.vertical_center.get_value() == pytest.approx(2.0)
+
+
+async def test_slits_set_reports_blade_progress_to_watchers(slits: Slits):
+    for motor in (slits.inboard, slits.outboard, slits.bottom, slits.top):
+        set_mock_value(motor.user_readback, 0.0)
+    updates = []
+
+    status = slits.set((10, 10, 4, 2))
+    status.watch(lambda **kw: updates.append(kw["name"]))
+    await status
+
+    assert {
+        slits.inboard.name,
+        slits.outboard.name,
+        slits.bottom.name,
+        slits.top.name,
+    } <= set(updates)