feat: Support for Device Composite injection (#1231)

This introduces support for composite devices in blueapi plan
signatures.

Composite devices must extend from pydantic `BaseModel`, and the devices
contained in them will be obtained from the `BlueskyContext`

---------

Co-authored-by: Dominic Oram <[email protected]>
Co-authored-by: Joseph Ware <[email protected]>

Author: 3 people

src/blueapi/core/context.py

@@ -2,7 +2,7 @@
 from collections.abc import Callable
 from dataclasses import InitVar, dataclass, field
 from importlib import import_module
-from inspect import Parameter, signature
+from inspect import Parameter, isclass, signature
 from types import ModuleType, NoneType, UnionType
 from typing import Any, Generic, TypeVar, Union, get_args, get_origin, get_type_hints
 
@@ -11,7 +11,7 @@
 from dodal.common.beamlines.beamline_utils import get_path_provider, set_path_provider
 from dodal.utils import AnyDevice, make_all_devices
 from ophyd_async.core import NotConnected
-from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, create_model
+from pydantic import BaseModel, GetCoreSchemaHandler, GetJsonSchemaHandler, create_model
 from pydantic.fields import FieldInfo
 from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema
 from pydantic_core import CoreSchema, core_schema
@@ -82,6 +82,9 @@ def is_bluesky_type(typ: type) -> bool:
     return typ in BLUESKY_PROTOCOLS or isinstance(typ, BLUESKY_PROTOCOLS)
 
 
+C = TypeVar("C", bound=BaseModel, covariant=True)
+
+
 @dataclass
 class BlueskyContext:
     """
@@ -383,7 +386,14 @@ def _type_spec_for_function(
                 )
 
             no_default = para.default is Parameter.empty
-            factory = None if no_default else DefaultFactory(para.default)
+            default_factory = (
+                self._composite_factory(arg_type)
+                if isclass(arg_type)
+                and issubclass(arg_type, BaseModel)
+                and isinstance(para.default, str)
+                else DefaultFactory(para.default)
+            )
+            factory = None if no_default else default_factory
             new_args[name] = (
                 self._convert_type(arg_type, no_default),
                 FieldInfo(default_factory=factory),
@@ -419,6 +429,20 @@ def _convert_type(self, typ: type | Any, no_default: bool = True) -> type:
             return root[new_types] if root else typ  # type: ignore
         return typ
 
+    def _composite_factory(self, composite_class: type[C]) -> Callable[[], C]:
+        def _inject_composite():
+            devices = {
+                field: self.find_device(info.default)
+                if info.annotation is not None
+                and is_bluesky_type(info.annotation)
+                and isinstance(info.default, str)
+                else info.default
+                for field, info in composite_class.model_fields.items()
+            }
+            return composite_class(**devices)
+
+        return _inject_composite
+
 
 D = TypeVar("D")
 

tests/unit_tests/worker/test_task_worker.py

@@ -21,6 +21,7 @@
 from blueapi.config import EnvironmentConfig, Source, SourceKind
 from blueapi.core import BlueskyContext, EventStream
 from blueapi.core.bluesky_types import DataEvent
+from blueapi.utils.base_model import BlueapiBaseModel
 from blueapi.worker import (
     Task,
     TaskStatus,
@@ -53,12 +54,9 @@
 class FakeDevice(Movable[float]):
     event: threading.Event
 
-    @property
-    def name(self) -> str:
-        return "fake_device"
-
-    def __init__(self) -> None:
+    def __init__(self, name: str = "fake_device") -> None:
         self.event = threading.Event()
+        self.name = name
 
     def set(self, value: float) -> Status:
         def when_done(_: Status):
@@ -84,14 +82,31 @@ def fake_device() -> FakeDevice:
 
 
 @pytest.fixture
-def context(fake_device: FakeDevice) -> BlueskyContext:
+def second_fake_device() -> FakeDevice:
+    return FakeDevice("second_fake_device")
+
+
+@pytest.fixture
+def context(fake_device: FakeDevice, second_fake_device: FakeDevice) -> BlueskyContext:
     ctx = BlueskyContext()
     ctx_config = EnvironmentConfig()
     ctx_config.sources.append(
         Source(kind=SourceKind.DEVICE_FUNCTIONS, module="devices")
     )
     ctx.register_plan(failing_plan)
     ctx.register_device(fake_device)
+    ctx.register_device(second_fake_device)
+    ctx.with_config(ctx_config)
+    return ctx
+
+
+@pytest.fixture
+def context_without_devices() -> BlueskyContext:
+    ctx = BlueskyContext()
+    ctx_config = EnvironmentConfig()
+    ctx_config.sources.append(
+        Source(kind=SourceKind.DEVICE_FUNCTIONS, module="devices")
+    )
     ctx.with_config(ctx_config)
     return ctx
 
@@ -681,3 +696,38 @@ def test_cycle_without_otel_context(mock_logger: Mock, inert_worker: TaskWorker)
     task.is_complete = False
     task.is_pending = True
     mock_logger.info.assert_called_with(f"Got new task: {task}")
+
+
+class MyComposite(BlueapiBaseModel):
+    dev_a: FakeDevice = inject(fake_device.name)
+    dev_b: FakeDevice = inject(second_fake_device.name)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+def injected_device_plan(composite: MyComposite = inject("")) -> MsgGenerator:
+    yield from ()
+
+
+def test_injected_composite_devices_are_found(
+    fake_device: FakeDevice,
+    second_fake_device: FakeDevice,
+    context: BlueskyContext,
+):
+    context.register_plan(injected_device_plan)
+    params = Task(name="injected_device_plan").prepare_params(context)
+    assert params["composite"].dev_a == fake_device
+    assert params["composite"].dev_b == second_fake_device
+
+
+def test_plan_module_with_composite_devices_can_be_loaded_before_device_module(
+    context_without_devices: BlueskyContext,
+    fake_device: FakeDevice,
+    second_fake_device: FakeDevice,
+):
+    context_without_devices.register_plan(injected_device_plan)
+    context_without_devices.register_device(fake_device)
+    context_without_devices.register_device(second_fake_device)
+    params = Task(name="injected_device_plan").prepare_params(context_without_devices)
+    assert params["composite"].dev_a == fake_device
+    assert params["composite"].dev_b == second_fake_device