Skip to content

feat: Support for dataclasses as composite types #1242

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
+152 −25
Open

feat: Support for dataclasses as composite types #1242

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1162236
Support for dataclasses as composite types
rtuck99 Oct 17, 2025
42486d5
Update to documentation
rtuck99 Oct 17, 2025
7bac9fb
Fix unit test build with example code modules
rtuck99 Oct 20, 2025
27be0b5
Make code coverage happy
rtuck99 Oct 20, 2025
015ecae
Fix plan schema for composites
rtuck99 Oct 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion docs/how-to/write-plans.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,24 @@ The type annotations in the example above (e.g. `: str`, `: int`, `-> MsgGenerat

## Injecting Devices

Some plans are created for specific sets of devices, or will almost always be used with the same devices, it is useful to be able to specify defaults. [Dodal makes this easy with its factory functions](https://diamondlightsource.github.io/dodal/main/how-to/include-devices-in-plans.html).
Some plans are created for specific sets of devices, or will almost always be used with the same devices, it is useful to be able to specify defaults. [Dodal makes this easy with its inject function](https://diamondlightsource.github.io/dodal/main/reference/generated/dodal.common.html#dodal.common.inject).

## Injecting multiple devices

If a plan requires multiple devices to be injected at once, rather than have a plan with several device parameters each of them with their own injection default, it is preferable to define a device composite which can be accepted as a parameter.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't agree that it is preferable- if you require overriding one device you must override the whole Composite.

I just think it's an alternative option that may be useful quite often, depending on how a specific beamline operates.

I'll look through the code change later, but I am admitting upfront that I am opinionated

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm. Perhaps I should rephrase that to reflect the difference between UDC and non-UDC use cases.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure it's a UDC vs. non-UDC thing. It's just if you have loads of devices it's cleaner to use a composite. If you have a couple it's cleaner not to

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think it's number but flexibility

class XYComposite(BaseModel):
  x: Motor
  y: Motor

def my_composite_plan(xy: XYComposite = inject()) -> MsgGenerator:
  ...


def my_individual_plan(x: Motor = inject(), y: Motor = inject()) -> MsgGenerator:
  ...

the json for the former if I wanted to scan theta/y instead of x/y:

{"xy": {"x": "theta", "y": "y"}}

compared to the latter:

{"x": "theta"}

The more devices I have in my BaseModel, the more verbose my plan signature would become if I don't use one, but the more duplicates I need to have in my JSON if I want to replace one of them.

If I'm always scanning x, y or I have multiple plans that all need x, y then a composite is neater. I'd prefer to define a specific device like a Stage to enforce a meaning to the relationship between x and y, but it's not always possible.

Or we support multiple levels of injection:

class MyComposite(BaseModel):
  x: Motor = inject()
  y: Motor = inject()

def my_plan(xy: MyComposite = inject()) -> MsgGenerator:  # uses x, y
  ...
{"xy": {"x": "theta"}}

Creating a new MyComposite with x=theta, y=inject("y") which finds y to pass in as the xy argument?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The way that I see it, if you are writing a plan which is intended to be consumed directly by scientists, for commissioning or one-off scripts explicitly inject the individual devices in the parameters, since that provides more flexibility over device selection. But I think most of those will probably not be called via blueapi.

There may be a use case with blueapi for individual device injection for things such as UI components if there are multiple similar devices e.g. slits, mirrors etc. and we want to provide direct control / monitoring over them for commissioning purposes. (although currently this functionality is mainly provided by the EDM)

If you have a plan that is intended to support a single frequently invoked high-level plan that will be called from a client application such as the hyperion supervisor, or GDA, then it's better to have a composite, because you almost certainly don't want or need the flexibility to inject different devices - client applications probably don't want to have to know about a beamline's specific configuration or what the naming of devices is, they just want to run a plan with whatever experimental parameters it requires.

If there's a choice of different devices for some section of the plan, the plan probably has to do other things that are bespoke for each device. For example, on i03, we can run with the panda or the zebra, but the initialisation and tear down code is different for each, so there would be no benefit to injecting only one device or the other.


For example you could define a composite as below:

```{literalinclude} ../../tests/unit_tests/code_examples/device_composite.py
:language: python
```

Then in your plan module:

```{literalinclude} ../../tests/unit_tests/code_examples/plan_with_composite.py
:language: python
```


## Injecting Metadata

Expand Down
50 changes: 32 additions & 18 deletions src/blueapi/core/context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging
from collections.abc import Callable
from dataclasses import InitVar, dataclass, field
from dataclasses import InitVar, dataclass, field, fields, is_dataclass
from importlib import import_module
from inspect import Parameter, isclass, signature
from types import ModuleType, NoneType, UnionType
Expand All @@ -11,7 +11,12 @@
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 BaseModel, 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
Expand Down Expand Up @@ -82,7 +87,7 @@ def is_bluesky_type(typ: type) -> bool:
return typ in BLUESKY_PROTOCOLS or isinstance(typ, BLUESKY_PROTOCOLS)


C = TypeVar("C", bound=BaseModel, covariant=True)
C = TypeVar("C", covariant=True)


@dataclass
Expand Down Expand Up @@ -386,16 +391,19 @@ def _type_spec_for_function(
)

no_default = para.default is Parameter.empty
default_factory = (
self._composite_factory(arg_type)
if isclass(arg_type)
and issubclass(arg_type, BaseModel)
if (
isclass(arg_type)
and (issubclass(arg_type, BaseModel) or is_dataclass(arg_type))
and isinstance(para.default, str)
else DefaultFactory(para.default)
)
):
default_factory = self._composite_factory(arg_type)
_type = SkipJsonSchema[self._convert_type(arg_type, no_default)]
else:
default_factory = DefaultFactory(para.default)
_type = self._convert_type(arg_type, no_default)
factory = None if no_default else default_factory
new_args[name] = (
self._convert_type(arg_type, no_default),
_type,
FieldInfo(default_factory=factory),
)
return new_args
Expand Down Expand Up @@ -431,14 +439,20 @@ def _convert_type(self, typ: type | Any, no_default: bool = True) -> type:

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()
}
if issubclass(composite_class, BaseModel):
devices = {
field_name: self.find_device(field_name)
for field_name in composite_class.model_fields.keys()
}
else:
assert is_dataclass(composite_class), (
f"Unsupported composite type: {composite_class}, composite must be"
" a pydantic BaseModel or a dataclass"
)
devices = {
field.name: self.find_device(field.name)
for field in fields(composite_class)
}
return composite_class(**devices)

return _inject_composite
Expand Down
8 changes: 8 additions & 0 deletions tests/unit_tests/code_examples/device_composite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import pydantic
from tests.unit_tests.code_examples.device_module import BimorphMirror


@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
class MyDeviceComposite:
oav: BimorphMirror
# More devices here....
12 changes: 12 additions & 0 deletions tests/unit_tests/code_examples/plan_with_composite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
from bluesky.utils import MsgGenerator
from dodal.common import inject
from tests.unit_tests.code_examples.device_composite import MyDeviceComposite


def my_plan(
parameter_one: int,
parameter_two: str,
my_necessary_devices: MyDeviceComposite = inject(""),
) -> MsgGenerator[None]:
# logic goes here
...
88 changes: 82 additions & 6 deletions tests/unit_tests/worker/test_task_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import dataclasses
import itertools
import threading
from collections.abc import Callable, Iterable
Expand All @@ -7,6 +8,7 @@
from typing import Any, TypeVar
from unittest.mock import ANY, MagicMock, Mock, patch

import pydantic
import pytest
from bluesky.protocols import Movable, Status
from bluesky.utils import MsgGenerator
Expand All @@ -21,6 +23,7 @@
from blueapi.config import EnvironmentConfig, Source, SourceKind
from blueapi.core import BlueskyContext, EventStream
from blueapi.core.bluesky_types import DataEvent
from blueapi.service.model import PlanModel
from blueapi.utils.base_model import BlueapiBaseModel
from blueapi.worker import (
Task,
Expand Down Expand Up @@ -661,6 +664,21 @@ def injected_device_plan(
assert params["dev"] == fake_device


def test_injected_devices_plan_model(
fake_device: FakeDevice,
context: BlueskyContext,
):
def injected_device_plan(
dev: FakeDevice = inject(fake_device.name),
) -> MsgGenerator:
yield from ()

context.register_plan(injected_device_plan)
plan = context.plans["injected_device_plan"]
model = PlanModel.from_plan(plan)
print(model)


def test_missing_injected_devices_fail_early(
context: BlueskyContext,
):
Expand Down Expand Up @@ -699,25 +717,83 @@ def test_cycle_without_otel_context(mock_logger: Mock, inert_worker: TaskWorker)


class MyComposite(BlueapiBaseModel):
dev_a: FakeDevice = inject(fake_device.name)
dev_b: FakeDevice = inject(second_fake_device.name)
fake_device: FakeDevice
second_fake_device: FakeDevice

model_config = {"arbitrary_types_allowed": True}


@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
class MyPydanticDataClassComposite:
fake_device: FakeDevice
second_fake_device: FakeDevice


@dataclasses.dataclass()
class MyStandardDataClassComposite:
fake_device: FakeDevice
second_fake_device: FakeDevice


def injected_device_plan(composite: MyComposite = inject("")) -> MsgGenerator:
yield from ()


def injected_dataclass_device_plan(
composite: MyPydanticDataClassComposite = inject(""),
) -> MsgGenerator:
yield from ()


def injected_standard_dataclass_device_plan(
composite: MyStandardDataClassComposite = 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
assert params["composite"].fake_device == fake_device
assert params["composite"].second_fake_device == second_fake_device


def test_injected_composite_devices_plan_model(
fake_device: FakeDevice,
second_fake_device: FakeDevice,
context: BlueskyContext,
):
context.register_plan(injected_device_plan)
plan = context.plans["injected_device_plan"]
PlanModel.from_plan(plan)


def test_injected_composite_with_pydantic_dataclass(
context: BlueskyContext,
fake_device: FakeDevice,
second_fake_device: FakeDevice,
):
context.register_plan(injected_dataclass_device_plan)
params = Task(name="injected_dataclass_device_plan").prepare_params(context)
assert params["composite"].fake_device == fake_device
assert params["composite"].second_fake_device == second_fake_device


def test_injected_composite_with_standard_dataclass(
context: BlueskyContext,
fake_device: FakeDevice,
second_fake_device: FakeDevice,
):
context.register_plan(injected_standard_dataclass_device_plan)
params = Task(name="injected_standard_dataclass_device_plan").prepare_params(
context
)
assert params["composite"].fake_device == fake_device
assert params["composite"].second_fake_device == second_fake_device


def test_plan_module_with_composite_devices_can_be_loaded_before_device_module(
Expand All @@ -729,5 +805,5 @@ def test_plan_module_with_composite_devices_can_be_loaded_before_device_module(
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
assert params["composite"].fake_device == fake_device
assert params["composite"].second_fake_device == second_fake_device