From 01092db9f36c569d422d1f401c056861174c94cd Mon Sep 17 00:00:00 2001 From: "shreelakshmi.iyengar@diamond.ac.uk" Date: Thu, 5 Nov 2026 17:33:50 +0000 Subject: [PATCH 1/3] change Plan to be a dataclass which has plan_function in it --- src/blueapi/core/bluesky_types.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index 64a2a5289e..2a96f473b0 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/src/blueapi/core/bluesky_types.py @@ -1,5 +1,6 @@ import inspect from collections.abc import Callable, Generator, Mapping +from dataclasses import dataclass from typing import ( Any, Protocol, @@ -25,7 +26,7 @@ ) from bluesky.utils import Msg, MsgGenerator from ophyd_async.core import Device as AsyncDevice -from pydantic import BaseModel, Field +from pydantic import BaseModel from blueapi.utils import BlueapiBaseModel @@ -81,18 +82,16 @@ def is_bluesky_plan_generator(func: PlanGenerator) -> bool: return False -class Plan(BlueapiBaseModel): +@dataclass +class Plan: """ A plan that can be run """ - name: str = Field(description="Referenceable name of the plan") - description: str | None = Field( - description="Description/docstring of the plan", default=None - ) - model: type[BaseModel] = Field( - description="Validation model of the parameters for the plan" - ) + name: str + model: type[BaseModel] + plan_function: PlanGenerator + description: str | None = None class DataEvent(BlueapiBaseModel): From 774e7c559d958b4b4d41531b7b22b2b3e53ddd6e Mon Sep 17 00:00:00 2001 From: "shreelakshmi.iyengar@diamond.ac.uk" Date: Thu, 5 Nov 2026 17:34:59 +0000 Subject: [PATCH 2/3] account for the change of adding plan_function to Plan --- src/blueapi/core/context.py | 7 ++++--- src/blueapi/worker/task.py | 5 ++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index f8a96b0d6a..c6738e8ca3 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -129,7 +129,6 @@ class BlueskyContext: path_provider: PathProvider | None = None plans: dict[str, Plan] = field(default_factory=dict) devices: dict[str, Device] = field(default_factory=dict) - plan_functions: dict[str, PlanGenerator] = field(default_factory=dict) _reference_cache: dict[type, type] = field(default_factory=dict) @@ -328,9 +327,11 @@ def my_plan(a: int, b: str): ) LOGGER.debug("Registering plan %s from %s", plan.__name__, plan.__module__) self.plans[plan.__name__] = Plan( - name=plan.__name__, model=model, description=plan.__doc__ + name=plan.__name__, + model=model, + plan_function=plan, + description=plan.__doc__, ) - self.plan_functions[plan.__name__] = plan return plan def register_device(self, device: Device, name: str | None = None) -> None: diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index d4cf0e809a..4f91bd14f1 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -62,9 +62,8 @@ def _inner_prepare_params( self, ctx: BlueskyContext ) -> tuple[list[Any], Mapping[str, Any]]: plan = ctx.plans[self.name] - func = ctx.plan_functions[self.name] - sig = signature(func) + sig = signature(plan.plan_function) bound = sig.bind_partial(*self.args, **self.params) # Only validate explicitly supplied arguments. This allows Pydantic's @@ -106,7 +105,7 @@ def do_task(self, ctx: BlueskyContext) -> None: f"Asked to run plan {self.name} with {self.params} and " f"metadata {self.metadata} for all runs" ) - plan = ctx.plan_functions[self.name] + plan = ctx.plans[self.name].plan_function prepared_args, prepared_kwargs = self.prepare_params(ctx) ctx.run_engine.md.update(self.metadata) result = ctx.run_engine(plan(*prepared_args, **prepared_kwargs)) From 34ce28525c0b314bd19a0db8076ae35f0a83c05f Mon Sep 17 00:00:00 2001 From: "shreelakshmi.iyengar@diamond.ac.uk" Date: Thu, 5 Nov 2026 17:48:54 +0000 Subject: [PATCH 3/3] tests: changes due to Plan now containing plan_function --- tests/unit_tests/cli/test_cli.py | 14 +++++++------- tests/unit_tests/service/test_rest_api.py | 14 ++++++-------- tests/unit_tests/test_example_code.py | 2 +- 3 files changed, 14 insertions(+), 16 deletions(-) diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py index e8458f0ac1..f858e6ac0f 100644 --- a/tests/unit_tests/cli/test_cli.py +++ b/tests/unit_tests/cli/test_cli.py @@ -40,7 +40,7 @@ ScratchRepository, generate_config_schema, ) -from blueapi.core.bluesky_types import DataEvent, Plan +from blueapi.core.bluesky_types import DataEvent from blueapi.service.model import ( DeviceModel, DeviceResponse, @@ -154,12 +154,13 @@ async def set(self, value: ComplexType): ... @responses.activate def test_get_plans(runner: CliRunner): - plan = Plan(name="my-plan", model=MyModel) response = responses.add( responses.GET, "http://localhost:8000/plans", - json=PlanResponse(plans=[PlanModel.from_plan(plan)]).model_dump(), + json=PlanResponse( + plans=[PlanModel(name="my-plan", schema=MyModel.model_json_schema())] + ).model_dump(), status=200, ) @@ -835,17 +836,16 @@ class ExtendedModel(BaseModel): def test_plan_output_formatting(): """Test for alternative plan output formats""" - plan = Plan( + plan = PlanModel( name="my-plan", description=dedent("""\ Summary of description Rest of description """), - model=ExtendedModel, + schema=ExtendedModel.model_json_schema(), ) - plans = PlanResponse(plans=[PlanModel.from_plan(plan)]) - + plans = PlanResponse(plans=[plan]) compact = dedent("""\ my-plan Summary of description diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py index 115ae5aae4..4315eecf61 100644 --- a/tests/unit_tests/service/test_rest_api.py +++ b/tests/unit_tests/service/test_rest_api.py @@ -22,7 +22,6 @@ OIDCConfig, RestConfig, ) -from blueapi.core.bluesky_types import Plan from blueapi.service import interface, main from blueapi.service.authorization import OpaUserClient, opa from blueapi.service.interface import ( @@ -53,7 +52,6 @@ class MockCountModel(BaseModel): ... -COUNT = Plan(name="count", model=MockCountModel) FAKE_INSTRUMENT_SESSION = "cm12345-1" @@ -158,8 +156,8 @@ def test_rest_config_with_cors_gets_plan( class MyModel(BaseModel): id: str - plan = Plan(name="my-plan", model=MyModel) - mock_runner.run.return_value = [PlanModel.from_plan(plan)] + plan = PlanModel(name="my-plan", schema=MyModel.model_json_schema()) + mock_runner.run.return_value = [plan] response_get = client_with_cors.get("/plans") assert response_get.status_code == status.HTTP_200_OK @@ -191,8 +189,8 @@ def test_get_plans(mock_runner: Mock, client: TestClient) -> None: class MyModel(BaseModel): id: str - plan = Plan(name="my-plan", model=MyModel) - mock_runner.run.return_value = [PlanModel.from_plan(plan)] + plan = PlanModel(name="my-plan", schema=MyModel.model_json_schema()) + mock_runner.run.return_value = [plan] response = client.get("/plans") @@ -217,8 +215,8 @@ def test_get_plan_by_name(mock_runner: Mock, client: TestClient) -> None: class MyModel(BaseModel): id: str - plan = Plan(name="my-plan", model=MyModel) - mock_runner.run.return_value = PlanModel.from_plan(plan) + plan = PlanModel(name="my-plan", schema=MyModel.model_json_schema()) + mock_runner.run.return_value = plan response = client.get("/plans/my-plan") diff --git a/tests/unit_tests/test_example_code.py b/tests/unit_tests/test_example_code.py index 0d16a871f7..de05093312 100644 --- a/tests/unit_tests/test_example_code.py +++ b/tests/unit_tests/test_example_code.py @@ -51,4 +51,4 @@ def test_invalid_plan_args_are_invalid(): module = importlib.import_module("invalid_plan_args") context.with_plan_module(module) with pytest.raises(TypeError): - context.plan_functions["demo"](**{"foo": 1, "bar": 2}) + context.plans["demo"].plan_function(**{"foo": 1, "bar": 2})