WIP as basically all tests failing

src/mx_bluesky/common/parameters/gridscan.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 
 from dodal.devices.aperturescatterguard import ApertureValue
+from dodal.devices.detector.det_dim_constants import EIGER2_X_9M_SIZE
+from dodal.devices.detector.detector import DetectorParams
 from dodal.devices.fast_grid_scan import (
     ZebraGridScanParams,
 )
+from dodal.utils import get_beamline_name
 from pydantic import Field, PrivateAttr
 from scanspec.core import Path as ScanPath
 from scanspec.specs import Line, Static
@@ -19,12 +22,15 @@
     XyzStarts,
 )
 from mx_bluesky.common.parameters.constants import (
+    DetectorParamConstants,
     GridscanParamConstants,
     HardwareConstants,
 )
 
 """Parameter models in this file are abstract. They should be inherited by a top-level model"""
 
+DETECTOR_SIZE_PER_BEAMLINE = {"i02-1": EIGER2_X_9M_SIZE}
+
 
 class GridCommon(
     DiffractionExperimentWithSample,
@@ -71,6 +77,37 @@ class SpecifiedThreeDGridScan(
     _set_stub_offsets: bool = PrivateAttr(default_factory=lambda: False)
     features: FeatureFlags | None = None
 
+    @property
+    def detector_params(self):
+        self.det_dist_to_beam_converter_path = (
+            self.det_dist_to_beam_converter_path
+            or DetectorParamConstants.BEAM_XY_LUT_PATH
+        )
+        optional_args = {}
+        if self.run_number:
+            optional_args["run_number"] = self.run_number
+        assert self.detector_distance_mm is not None, (
+            "Detector distance must be filled before generating DetectorParams"
+        )
+        return DetectorParams(
+            detector_size_constants=DETECTOR_SIZE_PER_BEAMLINE[
+                get_beamline_name("i03")
+            ],
+            expected_energy_ev=self.demand_energy_ev,
+            exposure_time=self.exposure_time_s,
+            directory=self.storage_directory,
+            prefix=self.file_name,
+            detector_distance=self.detector_distance_mm,
+            omega_start=self.omega_start_deg or 0,
+            omega_increment=0,
+            num_images_per_trigger=1,
+            num_triggers=self.num_images,
+            use_roi_mode=self.use_roi_mode,
+            det_dist_to_beam_converter_path=self.det_dist_to_beam_converter_path,
+            trigger_mode=self.trigger_mode,
+            **optional_args,
+        )
+
     @property
     def FGS_params(self) -> ZebraGridScanParams:
         return ZebraGridScanParams(

src/mx_bluesky/common/plans/common_flyscan_xray_centre_plan.py

@@ -31,7 +31,6 @@
     ZOCALO_READING_PLAN_NAME,
     ZOCALO_STAGE_GROUP,
     XrcResult,
-    ZocaloResults,
     get_full_processing_results,
 )
 
@@ -50,14 +49,14 @@
 )
 from mx_bluesky.common.parameters.gridscan import SpecifiedThreeDGridScan
 from mx_bluesky.common.plans.do_fgs import kickoff_and_complete_gridscan
+from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.exceptions import (
     CrystalNotFoundException,
     SampleException,
 )
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.utils.tracing import TRACER
 from mx_bluesky.common.xrc_result import XRayCentreEventHandler, XRayCentreResult
-from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
 # Saves needing to write 'assert not None' everywhere
@@ -88,33 +87,47 @@ class ReadHardwareTime(StrEnum):
     PRE_COLLECTION = "pre collection"
 
 
+# TODO: Think about proper typing for all these
 @dataclasses.dataclass
 class BeamlineSpecificFGSFeatures(Generic[D, P]):
-    setup_trigger: Callable[[D, P], MsgGenerator]
+    setup_trigger_plan: Callable[[D, P], MsgGenerator]
     tidy_plan: Callable[[D], MsgGenerator]
-    set_flyscan_params: Callable[[], MsgGenerator]
+    set_flyscan_params_plan: Callable[[], MsgGenerator]
     fgs_motors: FastGridScanCommon
-    read_pre_collection_plan: Callable[[], MsgGenerator]
+    read_pre_flyscan_plan: Callable[[], MsgGenerator]
     read_during_collection_plan: Callable[[], MsgGenerator]
-    plan_using_xrc_results: Callable[[D, P, XRayCentreResult], MsgGenerator] = null_plan
+    plan_after_getting_xrc_results: Callable[..., MsgGenerator] = null_plan
 
 
 def construct_beamline_specific_FGS_features(
-    self,
-    setup_trigger: Callable[[D, P], MsgGenerator],
+    setup_trigger_plan: Callable[[D, P], MsgGenerator],
     tidy_plan: Callable[[D], MsgGenerator],
-    set_flyscan_params: Callable[[], MsgGenerator],
+    set_flyscan_params_plan: Callable[[], MsgGenerator],
     fgs_motors: FastGridScanCommon,
-    signals_to_read_pre_collection: list[Readable],
+    signals_to_read_pre_flyscan: list[Readable],
     signals_to_read_during_collection: list[Readable],
-    plan_using_xrc_results: Callable[
-        [D, P, XRayCentreResult], MsgGenerator
-    ] = null_plan,
+    plan_after_getting_xrc_results=null_plan,
 ) -> BeamlineSpecificFGSFeatures:
-    """Construct the class needed to do beamline-specific parts of the XRC FGS"""
-    read_pre_collection_plan = partial(
+    """Construct the class needed to do beamline-specific parts of the XRC FGS
+
+    Args:
+        setup_trigger_plan: Configure any triggering, for example with the Zebra or PandA device. Ran directly before kicking off the gridscan.
+
+        tidy_plan: Tidy up states of devices. Ran at the end of the flyscan, regardless of whether or not it finished successfully.
+
+        set_flyscan_params_plan: TODO this one and Make better type hints for it
+
+        fgs_motors: Composite device representing the fast grid scan's motion program parameters.
+
+        signals_to_read_pre_flyscan: Signals which will be read and saved as a bluesky event document after all configuration, but before the gridscan.
+
+        signals_to_read_during_collection: Signals which will be read and saved as a bluesky event document whilst the gridscan motion is in progress
+
+        plan_after_getting_xrc_results: Optional plan which is ran after x-ray centring results have been retrieved from Zocalo.
+    """
+    read_pre_flyscan_plan = partial(
         read_hardware_plan,
-        signals_to_read_pre_collection,
+        signals_to_read_pre_flyscan,
         ReadHardwareTime.DURING_COLLECTION,
     )
 
@@ -124,13 +137,13 @@ def construct_beamline_specific_FGS_features(
         ReadHardwareTime.DURING_COLLECTION,
     )
     return BeamlineSpecificFGSFeatures(
-        setup_trigger,
+        setup_trigger_plan,
         tidy_plan,
-        set_flyscan_params,
+        set_flyscan_params_plan,
         fgs_motors,
-        read_pre_collection_plan,
+        read_pre_flyscan_plan,
         read_during_collection_plan,
-        plan_using_xrc_results,
+        plan_after_getting_xrc_results,
     )
 
 
@@ -148,6 +161,7 @@ def read_hardware_plan(
     yield from bps.create(name=event_name)
     for signal in signals:
         yield from bps.read(signal)
+    yield from bps.save()
 
 
 def create_devices(context: BlueskyContext) -> FlyScanEssentialDevices:
@@ -177,8 +191,8 @@ def flyscan_and_fetch_results() -> MsgGenerator:
     assert xray_centre_results, (
         "Flyscan result event not received or no crystal found and exception not raised"
     )
-    # Typing complains if you don't put parameters here
-    yield from feature_controlled.plan_using_xrc_results(
+
+    yield from feature_controlled.plan_after_getting_xrc_results(
         composite, parameters, xray_centre_results[0]
     )
 
@@ -237,7 +251,7 @@ def run_gridscan_and_fetch_results(
     """A multi-run plan which runs a gridscan, gets the results from zocalo
     and fires an event with the centres of mass determined by zocalo"""
 
-    yield from feature_controlled.setup_trigger(fgs_composite, parameters)
+    yield from feature_controlled.setup_trigger_plan(fgs_composite, parameters)
 
     LOGGER.info("Starting grid scan")
     yield from bps.stage(
@@ -343,10 +357,10 @@ def run_gridscan(
     # we should generate an event reading the values which need to be included in the
     # ispyb deposition
     with TRACER.start_span("ispyb_hardware_readings"):
-        yield from feature_controlled.read_pre_collection_plan()
+        yield from feature_controlled.read_pre_flyscan_plan()
 
     LOGGER.info("Setting fgs params")
-    yield from feature_controlled.set_flyscan_params()
+    yield from feature_controlled.set_flyscan_params_plan()
 
     LOGGER.info("Waiting for gridscan validity check")
     yield from wait_for_gridscan_valid(feature_controlled.fgs_motors)

src/mx_bluesky/common/utils/context.py

@@ -0,0 +1,68 @@
+import dataclasses
+from typing import Any, ClassVar, Protocol, TypeVar, get_type_hints
+
+from blueapi.core import BlueskyContext
+from blueapi.core.bluesky_types import Device
+
+from mx_bluesky.common.utils.log import LOGGER
+
+T = TypeVar("T", bound=Device)
+
+
+class _IsDataclass(Protocol):
+    """Protocol followed by any dataclass"""
+
+    __dataclass_fields__: ClassVar[dict]
+
+
+DT = TypeVar("DT", bound=_IsDataclass)
+
+
+def find_device_in_context(
+    context: BlueskyContext,
+    name: str,
+    # Typing in here is wrong (see https://github.com/microsoft/pyright/issues/7228#issuecomment-1934500232)
+    # but this whole thing will go away when we do https://github.com/DiamondLightSource/hyperion/issues/868
+    expected_type: type[T] = Device,  # type: ignore
+) -> T:
+    LOGGER.debug(f"Looking for device {name} of type {expected_type} in context")
+
+    device = context.find_device(name)
+    if device is None:
+        raise ValueError(
+            f"Cannot find device named '{name}' in bluesky context {context.devices}."
+        )
+
+    if not isinstance(device, expected_type):
+        raise ValueError(
+            f"Found device named '{name}' and expected it to be a '{expected_type}' but it was a '{device.__class__.__name__}'"
+        )
+
+    LOGGER.debug(f"Found matching device {device}")
+    return device
+
+
+def device_composite_from_context(context: BlueskyContext, dc: type[DT]) -> DT:
+    """
+    Initializes all of the devices referenced in a given dataclass from a provided
+    context, checking that the types of devices returned by the context are compatible
+    with the type annotations of the dataclass.
+
+    Note that if the context was not created with `wait_for_connection=True` devices may
+    still be unconnected.
+    """
+    LOGGER.debug(
+        f"Attempting to initialize devices referenced in dataclass {dc} from blueapi context"
+    )
+
+    devices: dict[str, Any] = {}
+    dc_type_hints: dict[str, Any] = get_type_hints(dc)
+
+    for field in dataclasses.fields(dc):
+        device = find_device_in_context(
+            context, field.name, expected_type=dc_type_hints.get(field.name, Device)
+        )
+
+        devices[field.name] = device
+
+    return dc(**devices)

src/mx_bluesky/hyperion/experiment_plans/__init__.py

@@ -4,7 +4,7 @@
 """
 
 from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
-    flyscan_xray_centre,
+    hyperion_flyscan_xray_centre,
 )
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     grid_detect_then_xray_centre,
@@ -24,7 +24,7 @@
 )
 
 __all__ = [
-    "flyscan_xray_centre",
+    "hyperion_flyscan_xray_centre",
     "grid_detect_then_xray_centre",
     "rotation_scan",
     "pin_tip_centre_then_xray_centre",

src/mx_bluesky/hyperion/experiment_plans/change_aperture_then_move_plan.py

@@ -2,25 +2,33 @@
 import bluesky.preprocessors as bpp
 import numpy
 from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
-from dodal.devices.smargon import Smargon, StubPosition
+from dodal.devices.smargon import StubPosition
 
 from mx_bluesky.common.utils.log import LOGGER
 from mx_bluesky.common.utils.tracing import TRACER
 from mx_bluesky.common.xrc_result import XRayCentreResult
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_x_y_z
+from mx_bluesky.hyperion.experiment_plans.device_composites import (
+    GridDetectThenXRayCentreComposite,
+    HyperionFlyScanXRayCentreComposite,
+)
 from mx_bluesky.hyperion.parameters.gridscan import HyperionSpecifiedThreeDGridScan
 
 
 def change_aperture_then_move_to_xtal(
     best_hit: XRayCentreResult,
-    smargon: Smargon,
-    aperture_scatterguard: ApertureScatterguard,
+    composite: HyperionFlyScanXRayCentreComposite
+    | GridDetectThenXRayCentreComposite,  # Need both types - Hyperion can enter this plan from grid detect or from flyscan depending on entry point
     parameters: HyperionSpecifiedThreeDGridScan | None = None,
 ):
     """For the given x-ray centring result,
     * Change the aperture so that the beam size is comparable to the crystal size
     * Centre on the centre-of-mass
     * Reset the stub offsets if specified by params"""
+
+    smargon = composite.smargon
+    aperture_scatterguard = composite.aperture_scatterguard
+
     if best_hit.bounding_box_mm is not None:
         bounding_box_size = numpy.abs(
             best_hit.bounding_box_mm[1] - best_hit.bounding_box_mm[0]