diff --git a/src/blueapi/cli/format.py b/src/blueapi/cli/format.py index 29ba01b73..24ebfd403 100644 --- a/src/blueapi/cli/format.py +++ b/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")), " ")) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index b0258141f..b5250db8e 100644 --- a/src/blueapi/client/client.py +++ b/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,7 +154,7 @@ 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 @@ -163,7 +162,8 @@ def get_plan(self, name: str) -> PlanModel: 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) diff --git a/src/blueapi/client/plan_cache.py b/src/blueapi/client/plan_cache.py index acf5b27ee..f7668c001 100644 --- a/src/blueapi/client/plan_cache.py +++ b/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. diff --git a/src/blueapi/client/rest.py b/src/blueapi/client/rest.py index f07f6b659..ea4220f5e 100644 --- a/src/blueapi/client/rest.py +++ b/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 diff --git a/src/blueapi/core/__init__.py b/src/blueapi/core/__init__.py index 192a72adf..1d2f2c53e 100644 --- a/src/blueapi/core/__init__.py +++ b/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", diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index 64a2a5289..e819a64f6 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/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): """ diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index f8a96b0d6..51869f48e 100644 --- a/src/blueapi/core/context.py +++ b/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 diff --git a/src/blueapi/service/__init__.py b/src/blueapi/service/__init__.py index ad8ccb1f6..22c7d5472 100644 --- a/src/blueapi/service/__init__.py +++ b/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"] diff --git a/src/blueapi/service/interface.py b/src/blueapi/service/interface.py index 7df474557..7b989e30f 100644 --- a/src/blueapi/service/interface.py +++ b/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]: diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 2ba310787..9417e6195 100644 --- a/src/blueapi/service/main.py +++ b/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) diff --git a/src/blueapi/service/model.py b/src/blueapi/service/model.py index 1b690d5c0..ff6bf3aca 100644 --- a/src/blueapi/service/model.py +++ b/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): diff --git a/tests/unit_tests/cli/test_cli.py b/tests/unit_tests/cli/test_cli.py index e8458f0ac..0dec88fbc 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 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 diff --git a/tests/unit_tests/service/test_rest_api.py b/tests/unit_tests/service/test_rest_api.py index 115ae5aae..e14190996 100644 --- a/tests/unit_tests/service/test_rest_api.py +++ b/tests/unit_tests/service/test_rest_api.py @@ -22,7 +22,7 @@ OIDCConfig, RestConfig, ) -from blueapi.core.bluesky_types import Plan +from blueapi.core.bluesky_types import BlueapiPlan from blueapi.service import interface, main from blueapi.service.authorization import OpaUserClient, opa from blueapi.service.interface import ( @@ -53,7 +53,7 @@ class MockCountModel(BaseModel): ... -COUNT = Plan(name="count", model=MockCountModel) +COUNT = BlueapiPlan(name="count", model=MockCountModel) FAKE_INSTRUMENT_SESSION = "cm12345-1" @@ -158,8 +158,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 = BlueapiPlan(name="my-plan", model=MyModel) + mock_runner.run.return_value = [PlanModel.model_validate(plan.to_model())] response_get = client_with_cors.get("/plans") assert response_get.status_code == status.HTTP_200_OK @@ -191,8 +191,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 = BlueapiPlan(name="my-plan", model=MyModel) + mock_runner.run.return_value = [PlanModel.model_validate(plan.to_model())] response = client.get("/plans") @@ -217,8 +217,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 = BlueapiPlan(name="my-plan", model=MyModel) + mock_runner.run.return_value = PlanModel.model_validate(plan.to_model()) response = client.get("/plans/my-plan") diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index 8f0a1269c..3cf719a27 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -740,7 +740,7 @@ def injected_device_plan( context.register_plan(injected_device_plan) plan = context.plans["injected_device_plan"] - PlanModel.from_plan(plan) + PlanModel.model_validate(plan.to_model()) def test_missing_injected_devices_fail_early( @@ -853,7 +853,7 @@ def test_injected_composite_devices_plan_model( ): context.register_plan(injected_device_plan) plan = context.plans["injected_device_plan"] - PlanModel.from_plan(plan) + PlanModel.model_validate(plan.to_model()) def test_injected_composite_with_pydantic_dataclass(