diff --git a/src/blueapi/client/plan_cache.py b/src/blueapi/client/plan_cache.py index 210cf873b..634dec3ed 100644 --- a/src/blueapi/client/plan_cache.py +++ b/src/blueapi/client/plan_cache.py @@ -1,4 +1,5 @@ import logging +from inspect import Parameter from typing import Any from blueapi.service.model import PlanModel, TaskParametersV2, TaskRequestV2 @@ -57,6 +58,7 @@ def __init__(self, name, model: PlanModel, client: ClientProtocol): self.model = model self._client = client self.__doc__ = model.description + self.parameter_kinds = model.parameter_kinds def __call__(self, *args, **kwargs) -> Any: """Execute the plan with the supplied arguments. @@ -109,23 +111,52 @@ def _build_args(self, *args: Any, **kwargs: Any) -> TaskParametersV2: args, kwargs, ) - properties = list(self.properties) + kinds = self.parameter_kinds - if len(args) > len(properties): - raise TypeError(f"{self.name} got too many arguments") + positional_parameters = [ + name + for name, kind in kinds.items() + if kind + in (Parameter.POSITIONAL_ONLY.name, Parameter.POSITIONAL_OR_KEYWORD.name) + ] - if extra := {k for k in kwargs if k not in properties}: - raise TypeError(f"{self.name} got unexpected arguments: {extra}") - - positional_names = properties[: len(args)] - - if duplicate := set(positional_names) & kwargs.keys(): - name = next(iter(duplicate)) - raise TypeError(f"{self.name} got multiple values for {name}") - - supplied = set(positional_names) | kwargs.keys() + var_positional = any( + kind == Parameter.VAR_POSITIONAL.name for kind in kinds.values() + ) + var_keyword = any(kind == Parameter.VAR_KEYWORD.name for kind in kinds.values()) + + if not var_positional and len(args) > len(positional_parameters): + raise TypeError( + f"{self.name}({args}) got too many arguments. " + f"Recieved {len(args)}. Expected {len(positional_parameters)}" + ) + + # Check positional/keyword collisions. + for name in positional_parameters[: len(args)]: + if name in kwargs: + raise TypeError(f"{self.name} got multiple values for {name}") + + # Keyword arguments must correspond to a named parameter, unless + # the function has **kwargs. + if not var_keyword: + unexpected = set(kwargs) - set(kinds) + if unexpected: + raise TypeError(f"{self.name} got unexpected arguments: {unexpected}") + + # A keyword-only parameter cannot be supplied positionally. + # Since positional_parameters excludes KEYWORD_ONLY, this is + # automatically handled above. + + supplied = set(kwargs) | set(positional_parameters[: len(args)]) + + required = { + name + for name in self.required + if kinds[name] + not in (Parameter.VAR_POSITIONAL.name, Parameter.VAR_KEYWORD.name) + } - if missing := set(self.required) - supplied: + if missing := required - supplied: raise TypeError(f"Missing argument(s) for {missing}") return TaskParametersV2(args=args, kwargs=kwargs) @@ -137,50 +168,60 @@ def __repr__(self) -> str: the plan's JSON schema. Short signatures are rendered on one line; longer signatures are formatted across multiple lines. """ - required = set(self.required) + kinds = self.model.parameter_kinds def _format_arg(name: str, info: dict[str, Any]) -> str: - typ = _pretty_type(info) - default = info.get("default") + typ = self.model.parameter_types[name] + kind = kinds[name] + + if kind == Parameter.VAR_POSITIONAL.name: + return f"*{name}: {typ}" + + if kind == Parameter.VAR_KEYWORD.name: + return f"**{name}: {typ}" - if name in required: + if name in self.required: return f"{name}: {typ}" - if default := info.get("default"): - return f"{name}: {typ} = {default!r}" - return f"{name}: {typ} | None = None" - args = [_format_arg(name, info) for name, info in self.properties.items()] - single_line = f"{self.name}({', '.join(args)})" + if "default" in info: + return f"{name}: {typ} = {info['default']!r}" - if len(single_line) <= _REPR_MAX_LENGTH and len(args) <= _REPR_MAX_ARGS_INLINE: - return single_line + return f"{name}: {typ} = None" - indent = " " - # Fall back to multiline if too many arguments or too long. - multiline_args = ",\n".join(f"{indent}{arg}" for arg in args) - return f"{self.name}(\n{multiline_args}\n)" + args: list[str] = [] + names = list(self.properties) + + has_var_positional = Parameter.VAR_POSITIONAL.name in kinds.values() + + for i, name in enumerate(names): + kind = kinds[name] + # Positional-only parameters need a "/" after the last one. + if kind == Parameter.POSITIONAL_ONLY.name: + args.append(_format_arg(name, self.properties[name])) -def _pretty_type(schema: dict[str, Any]) -> str: - """Convert a JSON schema type definition into a readable Python type. + if ( + i + 1 == len(names) + or kinds[names[i + 1]] != Parameter.POSITIONAL_ONLY.name + ): + args.append("/") - Handles references, arrays, unions, and primitive JSON schema types. - Unknown or unsupported schemas fall back to ``Any`` where no useful type - information is available. - """ - if "$ref" in schema: - return schema["$ref"].split("/")[-1] + continue + + # Keyword-only parameters need a "*" separator if there is + # no *args parameter to provide the separator. + if kind == Parameter.KEYWORD_ONLY.name and not has_var_positional: + if not args or args[-1] != "*": + args.append("*") - if schema.get("type") == "array": - item_schema = schema.get("items", {}) - inner = _pretty_type(item_schema) - return f"list[{inner}]" + args.append(_format_arg(name, self.properties[name])) - if "anyOf" in schema: - return " | ".join(_pretty_type(s) for s in schema["anyOf"]) + single_line = f"{self.name}({', '.join(args)})" - json_type = schema.get("type") - if isinstance(json_type, str): - return _JSON_TYPE_MAP.get(json_type, json_type.split(".")[-1]) + if len(single_line) <= _REPR_MAX_LENGTH and len(args) <= _REPR_MAX_ARGS_INLINE: + return single_line - return "Any" + indent = " " + # Fall back to multiline if too many arguments or too long. + multiline_args = ",\n".join(f"{indent}{arg}" for arg in args) + return f"{self.name}(\n{multiline_args}\n)" diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index 64a2a5289..c3fb4be4f 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/src/blueapi/core/bluesky_types.py @@ -93,6 +93,8 @@ class Plan(BlueapiBaseModel): model: type[BaseModel] = Field( description="Validation model of the parameters for the plan" ) + parameter_kinds: dict[str, str] = {} + parameter_types: dict[str, Any] = {} class DataEvent(BlueapiBaseModel): diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index c59a70b67..76df2bb7a 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -326,9 +326,21 @@ def my_plan(a: int, b: str): __config__=BlueapiPlanModelConfig, **self._type_spec_for_function(plan), # type: ignore ) + parameter_kinds = { + name: parameter.kind.name + for name, parameter in signature(plan).parameters.items() + } + parameter_types = { + name: parameter.annotation + for name, parameter in signature(plan).parameters.items() + } 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, + description=plan.__doc__, + parameter_kinds=parameter_kinds, + parameter_types=parameter_types, ) self.plan_functions[plan.__name__] = plan return plan diff --git a/src/blueapi/service/model.py b/src/blueapi/service/model.py index b29b943ad..e69e84ffa 100644 --- a/src/blueapi/service/model.py +++ b/src/blueapi/service/model.py @@ -1,7 +1,8 @@ import uuid from collections.abc import Iterable from enum import StrEnum -from typing import Annotated, Any +from types import NoneType, UnionType +from typing import Annotated, Any, Union, get_args, get_origin from bluesky.protocols import HasName from pydantic import Field @@ -114,6 +115,36 @@ class DeviceResponse(BlueapiBaseModel): devices: list[DeviceModel] = Field(description="Devices available to use in plans") +def _pretty_annotation(annotation: Any) -> str: + # Unwrap Annotated[T, ...] + if get_origin(annotation) is Annotated: + annotation = get_args(annotation)[0] + + # None + if annotation is NoneType: + return "None" + + origin = get_origin(annotation) + args = get_args(annotation) + + # PEP 604: T | U + if origin is UnionType: + return " | ".join(_pretty_annotation(arg) for arg in args) + + # typing.Union[T, U] + if origin is Union: + return " | ".join(_pretty_annotation(arg) for arg in args) + + # Generic types, e.g. Movable[float], tuple[T, ...], list[T] + if origin is not None: + origin_name = getattr(origin, "__name__", str(origin)) + formatted_args = ", ".join(_pretty_annotation(arg) for arg in args) + return f"{origin_name}[{formatted_args}]" + + # Plain classes + return getattr(annotation, "__name__", str(annotation)) + + class PlanModel(BlueapiBaseModel): """ Representation of a plan @@ -128,6 +159,8 @@ class PlanModel(BlueapiBaseModel): alias="schema", default_factory=dict, ) + parameter_kinds: dict[str, str] = Field(default_factory=dict) + parameter_types: dict[str, str] = Field(default_factory=dict) @classmethod def from_plan(cls, plan: Plan) -> "PlanModel": @@ -135,6 +168,11 @@ def from_plan(cls, plan: Plan) -> "PlanModel": name=plan.name, schema=plan.model.model_json_schema(), description=plan.description, + parameter_kinds=plan.parameter_kinds, + parameter_types={ + name: _pretty_annotation(annotation) + for name, annotation in plan.parameter_types.items() + }, ) diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index 563f45f36..312d9d28f 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -90,7 +90,14 @@ def prepare_params( case Parameter.VAR_KEYWORD: kwargs.update(value) else: - # Let the generated Pydantic model provide the default. + # Defaults need to be materialised for ordinary parameters, + # but variadic parameters are absent when not supplied. + if parameter.kind in ( + Parameter.VAR_POSITIONAL, + Parameter.VAR_KEYWORD, + ): + continue + value = getattr(validated, name) kwargs[name] = value diff --git a/tests/unit_tests/client/test_plan_cache.py b/tests/unit_tests/client/test_plan_cache.py index 81fca4146..5676f7028 100644 --- a/tests/unit_tests/client/test_plan_cache.py +++ b/tests/unit_tests/client/test_plan_cache.py @@ -1,3 +1,4 @@ +from inspect import Parameter from textwrap import dedent from unittest.mock import Mock @@ -19,6 +20,11 @@ }, "required": ["one"], }, + # To Do, parameter kinds should probably be part of schema + parameter_kinds={ + "one": Parameter.POSITIONAL_OR_KEYWORD.name, + "two": Parameter.POSITIONAL_OR_KEYWORD.name, + }, )