refactor: combine Plan and PlanModel into BlueapiPlan

src/blueapi/cli/format.py

@@ -14,7 +14,6 @@
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
-    PlanModel,
     PlanResponse,
     PythonEnvironmentResponse,
     SourceInfo,
@@ -57,7 +56,8 @@ def display_full(obj: Any, stream: Stream):
         case PlanResponse(plans=plans):
             for plan in plans:
                 display_full(plan, stream)
-        case PlanModel(name=name, description=desc, parameter_schema=schema):
+        case {"name": name, "schema": schema, **rest}:
+            desc = rest.get("description")
             print(name)
             if desc:
                 print(indent(dedent(desc).strip(), "    "))
@@ -110,7 +110,15 @@ def display_json(obj: Any, stream: Stream):
         case BaseModel():
             print(json.dumps(obj.model_dump()))
         case list():
-            print(json.dumps([it.model_dump() for it in obj], indent=2))
+            print(
+                json.dumps(
+                    [
+                        it.model_dump() if isinstance(it, BaseModel) else it
+                        for it in obj
+                    ],
+                    indent=2,
+                )
+            )
         case _:
             print(json.dumps(obj))
 
@@ -121,7 +129,8 @@ def display_compact(obj: Any, stream: Stream):
         case PlanResponse(plans=plans):
             for plan in plans:
                 display_compact(plan, stream)
-        case PlanModel(name=name, description=desc, parameter_schema=schema):
+        case {"name": name, "schema": schema, **rest}:
+            desc = rest.get("description")
             print(name)
             if desc:
                 print(indent(dedent(desc.split("\n\n")[0].strip("\n")), "    "))

src/blueapi/client/client.py

@@ -6,7 +6,7 @@
 from contextlib import suppress
 from functools import cached_property
 from pathlib import Path
-from typing import Self
+from typing import Any, Self
 
 from bluesky_stomp.messaging import MessageContext, StompClient
 from bluesky_stomp.models import Broker
@@ -26,7 +26,6 @@
     DeviceResponse,
     EnvironmentResponse,
     OIDCConfig,
-    PlanModel,
     PlanResponse,
     PythonEnvironmentResponse,
     SourceInfo,
@@ -155,15 +154,16 @@ def get_plans(self) -> PlanResponse:
 
     @start_as_current_span(TRACER, "name")
     @deprecated("plans[name]")
-    def get_plan(self, name: str) -> PlanModel:
+    def get_plan(self, name: str) -> dict[str, Any]:
         """
         Get details of a single plan
 
         Args:
             name: Plan name
 
         Returns:
-            PlanModel: Details of the plan if found
+            dict[str, Any]: Plan definition returned by BlueapiPlan.to_model()
+            containing the plan name, description and parameter schema.
         """
         return self._rest.get_plan(name)
 

src/blueapi/client/plan_cache.py

@@ -2,7 +2,7 @@
 from typing import Any
 
 from blueapi.client.rest import InvalidParametersError
-from blueapi.service.model import PlanModel, TaskRequest
+from blueapi.service.model import TaskRequest
 from blueapi.worker.event import TaskError, TaskResult, TaskStatus
 
 from .protocols import ClientProtocol
@@ -29,10 +29,12 @@ def __init__(self, typ: str, message: str):
 class PlanCache:
     """Collection of plans that can be accessed by name or attribute."""
 
-    def __init__(self, client: ClientProtocol, plans: list[PlanModel]):
+    def __init__(self, client: ClientProtocol, plans: list[dict[str, Any]]):
         self._cache = {
-            model.name: Plan(name=model.name, model=model, client=client)
-            for model in plans
+            plan_dict["name"]: Plan(
+                name=plan_dict["name"], model=plan_dict, client=client
+            )
+            for plan_dict in plans
         }
         for name, plan in self._cache.items():
             if name.startswith("_"):
@@ -55,11 +57,11 @@ def __repr__(self) -> str:
 class Plan:
     """Callable client-side reference to a registered BlueAPI plan."""
 
-    def __init__(self, name: str, model: PlanModel, client: ClientProtocol):
+    def __init__(self, name: str, model: dict[str, Any], client: ClientProtocol):
         self.name = name
         self.model = model
         self._client = client
-        self.__doc__ = model.description
+        self.__doc__ = model.get("description")
 
     def __call__(self, *args, **kwargs) -> Any:
         """Execute the plan with the supplied arguments.
@@ -92,15 +94,15 @@ def __call__(self, *args, **kwargs) -> Any:
 
     @property
     def help_text(self) -> str:
-        return self.model.description or f"Plan {self!r}"
+        return self.model.get("description") or f"Plan {self!r}"
 
     @property
     def properties(self) -> dict[str, Any]:
-        return self.model.parameter_schema.get("properties", {})
+        return self.model.get("schema", {}).get("properties", {})
 
     @property
     def required(self) -> list[str]:
-        return self.model.parameter_schema.get("required", [])
+        return self.model.get("schema", {}).get("required", [])
 
     def __repr__(self) -> str:
         """Return a signature-like representation of the plan.

src/blueapi/client/rest.py

@@ -25,7 +25,6 @@
     DeviceResponse,
     EnvironmentResponse,
     OIDCConfig,
-    PlanModel,
     PlanResponse,
     PythonEnvironmentResponse,
     SourceInfo,
@@ -218,9 +217,9 @@ def __init__(
     def get_plans(self) -> PlanResponse:
         return self._request_and_deserialize("/plans", PlanResponse)
 
-    def get_plan(self, name: str) -> PlanModel:
+    def get_plan(self, name: str) -> dict[str, Any]:
         try:
-            return self._request_and_deserialize(f"/plans/{name}", PlanModel)
+            return self._request_and_deserialize(f"/plans/{name}", dict[str, Any])
         except NotFoundError as e:
             raise UnknownPlanError(404, f"Plan '{name}' not found") from e
 

src/blueapi/core/__init__.py

@@ -6,7 +6,6 @@
     DataEvent,
     Device,
     MsgGenerator,
-    Plan,
     PlanGenerator,
     WatchableStatus,
     is_bluesky_compatible_device,
@@ -20,7 +19,6 @@
 OTLP_EXPORT_ENABLED = environ.get("OTLP_EXPORT_ENABLED") == "true"
 
 __all__ = [
-    "Plan",
     "PlanGenerator",
     "MsgGenerator",
     "Device",

src/blueapi/core/bluesky_types.py

@@ -81,7 +81,7 @@ def is_bluesky_plan_generator(func: PlanGenerator) -> bool:
         return False
 
 
-class Plan(BlueapiBaseModel):
+class BlueapiPlan(BlueapiBaseModel):
     """
     A plan that can be run
     """
@@ -94,6 +94,13 @@ class Plan(BlueapiBaseModel):
         description="Validation model of the parameters for the plan"
     )
 
+    def to_model(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "schema": self.model.model_json_schema(),
+            "description": self.description,
+        }
+
 
 class DataEvent(BlueapiBaseModel):
     """

src/blueapi/core/context.py

@@ -42,8 +42,8 @@
 from .bluesky_types import (
     BLUESKY_PROTOCOLS,
     AsyncDevice,
+    BlueapiPlan,
     Device,
-    Plan,
     PlanGenerator,
     is_bluesky_compatible_device,
     is_bluesky_plan_generator,
@@ -127,7 +127,7 @@ class BlueskyContext:
     tiled_conf: TiledConfig | None = field(default=None, init=False, repr=False)
     numtracker: NumtrackerClient | None = field(default=None, init=False, repr=False)
     path_provider: PathProvider | None = None
-    plans: dict[str, Plan] = field(default_factory=dict)
+    plans: dict[str, BlueapiPlan] = field(default_factory=dict)
     devices: dict[str, Device] = field(default_factory=dict)
     plan_functions: dict[str, PlanGenerator] = field(default_factory=dict)
 
@@ -327,7 +327,7 @@ def my_plan(a: int, b: str):
             **self._type_spec_for_function(plan),  # type: ignore
         )
         LOGGER.debug("Registering plan %s from %s", plan.__name__, plan.__module__)
-        self.plans[plan.__name__] = Plan(
+        self.plans[plan.__name__] = BlueapiPlan(
             name=plan.__name__, model=model, description=plan.__doc__
         )
         self.plan_functions[plan.__name__] = plan

src/blueapi/service/__init__.py

@@ -1,4 +1,4 @@
 from .authentication import SessionManager
-from .model import DeviceModel, PlanModel
+from .model import DeviceModel
 
-__all__ = ["PlanModel", "DeviceModel", "SessionManager"]
+__all__ = ["DeviceModel", "SessionManager"]

src/blueapi/service/interface.py

@@ -22,7 +22,6 @@
 from blueapi.service.authentication import TiledAuth
 from blueapi.service.model import (
     DeviceModel,
-    PlanModel,
     PythonEnvironmentResponse,
     SourceInfo,
     TaskRequest,
@@ -179,14 +178,14 @@ def forward_message(event: Any, correlation_id: str | None) -> None:
     stream.subscribe(forward_message)
 
 
-def get_plans() -> list[PlanModel]:
+def get_plans() -> list[dict[str, Any]]:
     """Get all available plans in the BlueskyContext"""
-    return [PlanModel.from_plan(plan) for plan in context().plans.values()]
+    return [plan.to_model() for plan in context().plans.values()]
 
 
-def get_plan(name: str) -> PlanModel:
+def get_plan(name: str) -> dict[str, Any]:
     """Get plan by name from the BlueskyContext"""
-    return PlanModel.from_plan(context().plans[name])
+    return context().plans[name].to_model()
 
 
 def get_devices() -> list[DeviceModel]:

src/blueapi/service/main.py

@@ -2,7 +2,7 @@
 import urllib.parse
 from collections.abc import Awaitable, Callable
 from contextlib import asynccontextmanager
-from typing import Annotated
+from typing import Annotated, Any
 
 import jwt
 from bluesky._vendor.super_state_machine.errors import TransitionError
@@ -73,7 +73,6 @@
     EnvironmentResponse,
     Health,
     HealthProbeResponse,
-    PlanModel,
     PlanResponse,
     PythonEnvironmentResponse,
     SourceInfo,
@@ -293,7 +292,7 @@ def get_plans(runner: Annotated[WorkerDispatcher, Depends(_runner)]) -> PlanResp
 @start_as_current_span(TRACER, "name")
 def get_plan_by_name(
     name: str, runner: Annotated[WorkerDispatcher, Depends(_runner)]
-) -> PlanModel:
+) -> dict[str, Any]:
     """Retrieve information about a plan by its (unique) name."""
     return runner.run(interface.get_plan, name)
 

src/blueapi/service/model.py

@@ -8,7 +8,7 @@
 from pydantic.json_schema import SkipJsonSchema
 
 from blueapi.config import OIDCConfig
-from blueapi.core import BLUESKY_PROTOCOLS, Device, Plan
+from blueapi.core import BLUESKY_PROTOCOLS, Device
 from blueapi.core.context import generic_bounds
 from blueapi.utils import BlueapiBaseModel
 from blueapi.worker import WorkerState
@@ -93,30 +93,6 @@ class DeviceResponse(BlueapiBaseModel):
     devices: list[DeviceModel] = Field(description="Devices available to use in plans")
 
 
-class PlanModel(BlueapiBaseModel):
-    """
-    Representation of a plan
-    """
-
-    name: str = Field(description="Name of the plan")
-    description: str | SkipJsonSchema[None] = Field(
-        description="Docstring of the plan", default=None
-    )
-    parameter_schema: dict[str, Any] = Field(
-        description="Schema of the plan's parameters",
-        alias="schema",
-        default_factory=dict,
-    )
-
-    @classmethod
-    def from_plan(cls, plan: Plan) -> "PlanModel":
-        return cls(
-            name=plan.name,
-            schema=plan.model.model_json_schema(),
-            description=plan.description,
-        )
-
-
 class PlanRequest(BlueapiBaseModel):
     """
     A query for plans
@@ -130,7 +106,9 @@ class PlanResponse(BlueapiBaseModel):
     Response to a query for plans
     """
 
-    plans: list[PlanModel] = Field(description="Plans available to use by a worker")
+    plans: list[dict[str, Any]] = Field(
+        description="Plans available to use by a worker"
+    )
 
 
 class TaskResponse(BlueapiBaseModel):

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 BlueapiPlan, DataEvent
 from blueapi.service.model import (
     DeviceModel,
     DeviceResponse,
@@ -154,12 +154,14 @@ async def set(self, value: ComplexType): ...
 
 @responses.activate
 def test_get_plans(runner: CliRunner):
-    plan = Plan(name="my-plan", model=MyModel)
+    plan = BlueapiPlan(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.model_validate(plan.to_model())]
+        ).model_dump(),
         status=200,
     )
 
@@ -835,7 +837,7 @@ class ExtendedModel(BaseModel):
 def test_plan_output_formatting():
     """Test for alternative plan output formats"""
 
-    plan = Plan(
+    plan = BlueapiPlan(
         name="my-plan",
         description=dedent("""\
             Summary of description
@@ -844,7 +846,7 @@ def test_plan_output_formatting():
             """),
         model=ExtendedModel,
     )
-    plans = PlanResponse(plans=[PlanModel.from_plan(plan)])
+    plans = PlanResponse(plans=[PlanModel.model_validate(plan.to_model())])
 
     compact = dedent("""\
                 my-plan