From 8d2666bb3d9a67d0c675d421ac4491b365f79041 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Thu, 10 Sep 2026 10:15:29 +0000 Subject: [PATCH 1/3] Update to preserve *args information and make useable by client --- src/blueapi/client/client.py | 92 +++++++++++++++++++++++-------- src/blueapi/core/bluesky_types.py | 1 + src/blueapi/core/context.py | 33 +++++++++-- src/blueapi/service/model.py | 6 +- src/blueapi/worker/task.py | 9 ++- 5 files changed, 110 insertions(+), 31 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 3d77a93daa..5ef4d203a8 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -156,6 +156,7 @@ def __init__(self, name, model: PlanModel, client: "BlueapiClient"): self.model = model self._client = client self.__doc__ = model.description + self.parameter_kinds = model.parameter_kinds def __call__(self, *args, **kwargs) -> Any: req = TaskRequest( @@ -182,48 +183,95 @@ def required(self) -> list[str]: return self.model.parameter_schema.get("required", []) def _build_args(self, *args, **kwargs) -> TaskParams: - log.info( - "Building args for %s, using %s and %s", - "[" + ",".join(self.properties) + "]", - args, - kwargs, - ) + kinds = self.parameter_kinds + + positional_parameters = [ + name + for name, kind in kinds.items() + if kind in ("POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD") + ] - properties = list(self.properties) + var_positional = any(kind == "VAR_POSITIONAL" for kind in kinds.values()) + var_keyword = any(kind == "VAR_KEYWORD" for kind in kinds.values()) - if len(args) > len(properties): + if not var_positional and len(args) > len(positional_parameters): raise TypeError(f"{self.name} got too many arguments") - if extra := {k for k in kwargs if k not in properties}: - raise TypeError(f"{self.name} got unexpected arguments: {extra}") + # 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}") - positional_names = properties[: len(args)] + # 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}") - if duplicate := set(positional_names) & kwargs.keys(): - name = next(iter(duplicate)) - raise TypeError(f"{self.name} got multiple values for {name}") + # A keyword-only parameter cannot be supplied positionally. + # Since positional_parameters excludes KEYWORD_ONLY, this is + # automatically handled above. - supplied = set(positional_names) | kwargs.keys() + supplied = set(kwargs) | set(positional_parameters[: len(args)]) - if missing := set(self.required) - supplied: + required = { + name + for name in self.required + if kinds[name] not in ("VAR_POSITIONAL", "VAR_KEYWORD") + } + + if missing := required - supplied: raise TypeError(f"Missing argument(s) for {missing}") return TaskParams(args=args, kwargs=kwargs) def __repr__(self) -> str: - 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") + kind = kinds[name] + + if kind == "VAR_POSITIONAL": + return f"*{name}: {typ}" - if name in required: + if kind == "VAR_KEYWORD": + return f"**{name}: {typ}" + + if name in self.required: return f"{name}: {typ}" - if default := info.get("default"): - return f"{name}: {typ} = {default!r}" + + if "default" in info: + return f"{name}: {typ} = {info['default']!r}" + return f"{name}: {typ} | None = None" - args = [_format_arg(name, info) for name, info in self.properties.items()] + args: list[str] = [] + names = list(self.properties) + + has_var_positional = "VAR_POSITIONAL" in kinds.values() + + for i, name in enumerate(names): + kind = kinds[name] + + # Positional-only parameters need a "/" after the last one. + if kind == "POSITIONAL_ONLY": + args.append(_format_arg(name, self.properties[name])) + + if i + 1 == len(names) or kinds[names[i + 1]] != "POSITIONAL_ONLY": + args.append("/") + + continue + + # Keyword-only parameters need a "*" separator if there is + # no *args parameter to provide the separator. + if kind == "KEYWORD_ONLY" and not has_var_positional: + if not args or args[-1] != "*": + args.append("*") + + args.append(_format_arg(name, self.properties[name])) + single_line = f"{self.name}({', '.join(args)})" if len(single_line) <= _REPR_MAX_LENGTH and len(args) <= _REPR_MAX_ARGS_INLINE: diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index 64a2a5289e..b53d44bda8 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/src/blueapi/core/bluesky_types.py @@ -93,6 +93,7 @@ class Plan(BlueapiBaseModel): model: type[BaseModel] = Field( description="Validation model of the parameters for the plan" ) + parameter_kinds: dict[str, str] = {} class DataEvent(BlueapiBaseModel): diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 76ce74e647..07b92346a5 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -326,9 +326,16 @@ 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() + } 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, ) self.plan_functions[plan.__name__] = plan return plan @@ -439,19 +446,31 @@ def _type_spec_for_function( Returns: Mapping of {name: (type, default)} to be used by pydantic for deserialising - function arguments + function arguments """ args = signature(func).parameters types = get_type_hints(func) new_args: dict[str, tuple[type, FieldInfo]] = {} + for name, para in args.items(): arg_type = types.get(name, Parameter.empty) + if arg_type is Parameter.empty: raise ValueError( f"Type annotation is required for '{name}' in '{func.__name__}'" ) - no_default = para.default is Parameter.empty + match para.kind: + case Parameter.VAR_POSITIONAL: + arg_type = tuple[arg_type, ...] + case Parameter.VAR_KEYWORD: + arg_type = dict[str, arg_type] + + no_default = para.default is Parameter.empty and para.kind not in ( + Parameter.VAR_POSITIONAL, + Parameter.VAR_KEYWORD, + ) + if ( isclass(arg_type) and (issubclass(arg_type, BaseModel) or is_dataclass(arg_type)) @@ -462,9 +481,15 @@ def _type_spec_for_function( info = FieldInfo(default_factory=default_factory) else: _type = self._convert_type(arg_type, no_default) + match para.default: case Parameter.empty: - info = FieldInfo(default_factory=None) + if para.kind is Parameter.VAR_POSITIONAL: + info = FieldInfo(default_factory=tuple) + elif para.kind is Parameter.VAR_KEYWORD: + info = FieldInfo(default_factory=dict) + else: + info = FieldInfo(default_factory=None) case None: info = FieldInfo(default_factory=lambda: None) case _: diff --git a/src/blueapi/service/model.py b/src/blueapi/service/model.py index 7d7aaabc1d..b8cd12de12 100644 --- a/src/blueapi/service/model.py +++ b/src/blueapi/service/model.py @@ -89,10 +89,6 @@ class DeviceResponse(BlueapiBaseModel): 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 @@ -102,6 +98,7 @@ class PlanModel(BlueapiBaseModel): alias="schema", default_factory=dict, ) + parameter_kinds: dict[str, str] = Field(default_factory=dict) @classmethod def from_plan(cls, plan: Plan) -> "PlanModel": @@ -109,6 +106,7 @@ def from_plan(cls, plan: Plan) -> "PlanModel": name=plan.name, schema=plan.model.model_json_schema(), description=plan.description, + parameter_kinds=plan.parameter_kinds, ) diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index a7291ce684..816392eb0c 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -82,7 +82,14 @@ def prepare_params( 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 From a08d85403b5553184b7b0be4e92cdd46e2b05a32 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Thu, 10 Sep 2026 10:39:21 +0000 Subject: [PATCH 2/3] Update to use Parameter rather than str --- src/blueapi/client/client.py | 33 +++++++++++++++++++-------------- 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 5ef4d203a8..d9c02427dd 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -5,15 +5,13 @@ from concurrent.futures import Future from contextlib import suppress from functools import cached_property +from inspect import Parameter from pathlib import Path from typing import Any, Self from bluesky_stomp.messaging import MessageContext, StompClient from bluesky_stomp.models import Broker -from observability_utils.tracing import ( - get_tracer, - start_as_current_span, -) +from observability_utils.tracing import get_tracer, start_as_current_span from blueapi.config import ( ApplicationConfig, @@ -188,11 +186,14 @@ def _build_args(self, *args, **kwargs) -> TaskParams: positional_parameters = [ name for name, kind in kinds.items() - if kind in ("POSITIONAL_ONLY", "POSITIONAL_OR_KEYWORD") + if kind + in (Parameter.POSITIONAL_ONLY.name, Parameter.POSITIONAL_OR_KEYWORD.name) ] - var_positional = any(kind == "VAR_POSITIONAL" for kind in kinds.values()) - var_keyword = any(kind == "VAR_KEYWORD" for kind in kinds.values()) + 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} got too many arguments") @@ -218,7 +219,8 @@ def _build_args(self, *args, **kwargs) -> TaskParams: required = { name for name in self.required - if kinds[name] not in ("VAR_POSITIONAL", "VAR_KEYWORD") + if kinds[name] + not in (Parameter.VAR_POSITIONAL.name, Parameter.VAR_KEYWORD.name) } if missing := required - supplied: @@ -233,10 +235,10 @@ def _format_arg(name: str, info: dict[str, Any]) -> str: typ = _pretty_type(info) kind = kinds[name] - if kind == "VAR_POSITIONAL": + if kind == Parameter.VAR_POSITIONAL.name: return f"*{name}: {typ}" - if kind == "VAR_KEYWORD": + if kind == Parameter.VAR_KEYWORD.name: return f"**{name}: {typ}" if name in self.required: @@ -250,23 +252,26 @@ def _format_arg(name: str, info: dict[str, Any]) -> str: args: list[str] = [] names = list(self.properties) - has_var_positional = "VAR_POSITIONAL" in kinds.values() + 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 == "POSITIONAL_ONLY": + if kind == Parameter.POSITIONAL_ONLY.name: args.append(_format_arg(name, self.properties[name])) - if i + 1 == len(names) or kinds[names[i + 1]] != "POSITIONAL_ONLY": + if ( + i + 1 == len(names) + or kinds[names[i + 1]] != Parameter.POSITIONAL_ONLY.name + ): args.append("/") continue # Keyword-only parameters need a "*" separator if there is # no *args parameter to provide the separator. - if kind == "KEYWORD_ONLY" and not has_var_positional: + if kind == Parameter.KEYWORD_ONLY.name and not has_var_positional: if not args or args[-1] != "*": args.append("*") From d5001868b76787a68dda85dd41f029bf37753ae2 Mon Sep 17 00:00:00 2001 From: Oli Wenman Date: Thu, 10 Sep 2026 12:33:39 +0000 Subject: [PATCH 3/3] Fix __repr__ representation of function on client --- src/blueapi/client/client.py | 30 ++---------------------- src/blueapi/core/bluesky_types.py | 1 + src/blueapi/core/context.py | 5 ++++ src/blueapi/service/model.py | 38 ++++++++++++++++++++++++++++++- 4 files changed, 45 insertions(+), 29 deletions(-) diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index d9c02427dd..8800d424d4 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -56,13 +56,6 @@ _REPR_MAX_LENGTH = 100 _REPR_MAX_ARGS_INLINE = 3 -_JSON_TYPE_MAP = { - "string": "str", - "integer": "int", - "boolean": "bool", - "number": "float", - "object": "dict", -} class MissingInstrumentSessionError(Exception): @@ -232,7 +225,7 @@ def __repr__(self) -> str: kinds = self.model.parameter_kinds def _format_arg(name: str, info: dict[str, Any]) -> str: - typ = _pretty_type(info) + typ = self.model.parameter_types[name] kind = kinds[name] if kind == Parameter.VAR_POSITIONAL.name: @@ -247,7 +240,7 @@ def _format_arg(name: str, info: dict[str, Any]) -> str: if "default" in info: return f"{name}: {typ} = {info['default']!r}" - return f"{name}: {typ} | None = None" + return f"{name}: {typ} = None" args: list[str] = [] names = list(self.properties) @@ -887,22 +880,3 @@ class PlanFailedError(Exception): def __init__(self, typ: str, message: str): super().__init__(message) self._type = typ - - -def _pretty_type(schema: dict[str, Any]) -> str: - if "$ref" in schema: - return schema["$ref"].split("/")[-1] - - if schema.get("type") == "array": - item_schema = schema.get("items", {}) - inner = _pretty_type(item_schema) - return f"list[{inner}]" - - if "anyOf" in schema: - return " | ".join(_pretty_type(s) for s in schema["anyOf"]) - - json_type = schema.get("type") - if isinstance(json_type, str): - return _JSON_TYPE_MAP.get(json_type, json_type.split(".")[-1]) - - return "Any" diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index b53d44bda8..c3fb4be4fd 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/src/blueapi/core/bluesky_types.py @@ -94,6 +94,7 @@ class Plan(BlueapiBaseModel): 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 07b92346a5..e3d564863b 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -330,12 +330,17 @@ def my_plan(a: int, b: str): 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__, 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 b8cd12de12..56565285d4 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 @@ -88,6 +89,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): name: str = Field(description="Name of the plan") description: str | SkipJsonSchema[None] = Field( @@ -99,6 +130,7 @@ class PlanModel(BlueapiBaseModel): 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": @@ -107,6 +139,10 @@ def from_plan(cls, plan: Plan) -> "PlanModel": 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() + }, )