Update client to support plans with *args and better __repr__

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

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

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

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,13 +159,20 @@ 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":
         return cls(
             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()
+            },
         )
 
 

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
 

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