refactor: Combining Plan and Plan functions

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):

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:

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))

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

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")
 

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})