feat: Improve BlueapiClient to add plan parameter type hints

src/blueapi/client/client.py

@@ -5,7 +5,6 @@
 from concurrent.futures import Future
 from contextlib import suppress
 from functools import cached_property
-from itertools import chain
 from pathlib import Path
 from typing import Any, Self
 
@@ -56,6 +55,16 @@
 
 log = logging.getLogger(__name__)
 
+_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):
     pass
@@ -164,8 +173,8 @@ def help_text(self) -> str:
         return self.model.description or f"Plan {self!r}"
 
     @property
-    def properties(self) -> set[str]:
-        return self.model.parameter_schema.get("properties", {}).keys()
+    def properties(self) -> dict[str, Any]:
+        return self.model.parameter_schema.get("properties", {})
 
     @property
     def required(self) -> list[str]:
@@ -201,10 +210,30 @@ def _build_args(self, *args, **kwargs):
             raise TypeError(f"Missing argument(s) for {missing}")
         return params
 
-    def __repr__(self):
-        opts = [p for p in self.properties if p not in self.required]
-        params = ", ".join(chain(self.required, (f"{opt}=None" for opt in opts)))
-        return f"{self.name}({params})"
+    def __repr__(self) -> str:
+        required = set(self.required)
+
+        def _format_arg(name: str, info: dict[str, Any]) -> str:
+            typ = _pretty_type(info)
+            default = info.get("default")
+
+            if name in required:
+                return f"{name}: {typ}"
+            if default := info.get("default"):
+                return f"{name}: {typ} = {default!r}"
+            return f"{name}: {typ} | None = None"
+
+        props: dict = self.model.parameter_schema.get("properties", {})
+        args = [_format_arg(name, info) for name, info in props.items()]
+        single_line = f"{self.name}({', '.join(args)})"
+
+        if len(single_line) <= _REPR_MAX_LENGTH and len(args) <= _REPR_MAX_ARGS_INLINE:
+            return single_line
+
+        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)"
 
 
 class BlueapiClient:
@@ -785,3 +814,22 @@ 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"

src/blueapi/core/context.py

@@ -5,7 +5,7 @@
 from importlib import import_module, metadata
 from inspect import Parameter, isclass, signature
 from types import ModuleType, NoneType, UnionType
-from typing import Any, Generic, TypeVar, Union, get_args, get_origin, get_type_hints
+from typing import Any, TypeVar, Union, get_args, get_origin, get_type_hints
 
 from bluesky.protocols import HasName
 from bluesky.run_engine import RunEngine
@@ -459,14 +459,19 @@ def _type_spec_for_function(
             ):
                 default_factory = self._composite_factory(arg_type)
                 _type = SkipJsonSchema[self._convert_type(arg_type, no_default)]
+                info = FieldInfo(default_factory=default_factory)
             else:
-                default_factory = DefaultFactory(para.default)
                 _type = self._convert_type(arg_type, no_default)
-            factory = None if no_default else default_factory
-            new_args[name] = (
-                _type,
-                FieldInfo(default_factory=factory),
-            )
+                match para.default:
+                    case Parameter.empty:
+                        info = FieldInfo(default_factory=None)
+                    case None:
+                        info = FieldInfo(default_factory=lambda: None)
+                    case _:
+                        info = FieldInfo(default=para.default)
+
+            new_args[name] = (_type, info)
+
         return new_args
 
     def _convert_type(self, typ: Any, no_default: bool = True) -> type:
@@ -517,19 +522,3 @@ def _inject_composite():
             return composite_class(**devices)
 
         return _inject_composite
-
-
-D = TypeVar("D")
-
-
-class DefaultFactory(Generic[D]):
-    _value: D
-
-    def __init__(self, value: D):
-        self._value = value
-
-    def __call__(self) -> D:
-        return self._value
-
-    def __eq__(self, other) -> bool:
-        return other.__class__ == self.__class__ and self._value == other._value

tests/system_tests/plans.json

@@ -29,6 +29,7 @@
                         "type": "array"
                     },
                     "num": {
+                        "default": 1,
                         "title": "Num",
                         "type": "integer"
                     },
@@ -44,6 +45,7 @@
                                 "type": "array"
                             }
                         ],
+                        "default": 0,
                         "title": "Delay"
                     },
                     "metadata": {
@@ -116,6 +118,7 @@
                         "type": "array"
                     },
                     "snake_axes": {
+                        "default": true,
                         "title": "Snake Axes",
                         "type": "boolean"
                     },
@@ -190,6 +193,7 @@
                         "type": "array"
                     },
                     "snake_axes": {
+                        "default": true,
                         "title": "Snake Axes",
                         "type": "boolean"
                     },
@@ -413,6 +417,7 @@
                                 "type": "boolean"
                             }
                         ],
+                        "default": true,
                         "title": "Snake Axes"
                     },
                     "metadata": {
@@ -495,6 +500,7 @@
                                 "type": "boolean"
                             }
                         ],
+                        "default": true,
                         "title": "Snake Axes"
                     },
                     "metadata": {
@@ -1360,6 +1366,7 @@
                         "type": "array"
                     },
                     "snake_axes": {
+                        "default": true,
                         "title": "Snake Axes",
                         "type": "boolean"
                     },
@@ -1434,6 +1441,7 @@
                         "type": "array"
                     },
                     "snake_axes": {
+                        "default": true,
                         "title": "Snake Axes",
                         "type": "boolean"
                     },
@@ -1613,6 +1621,7 @@
                         "type": "string"
                     },
                     "wait": {
+                        "default": true,
                         "title": "Wait",
                         "type": "boolean"
                     }
@@ -1647,6 +1656,7 @@
                         "type": "string"
                     },
                     "wait": {
+                        "default": true,
                         "title": "Wait",
                         "type": "boolean"
                     }

tests/unit_tests/client/test_client.py

@@ -1,5 +1,6 @@
 import uuid
 from collections.abc import Callable
+from textwrap import dedent
 from unittest.mock import MagicMock, Mock, call, patch
 
 import pytest
@@ -806,7 +807,71 @@ def test_plan_fallback_help_text(client):
         ),
         client,
     )
-    assert plan.help_text == "Plan foo(one, two=None)"
+    assert plan.help_text == "Plan foo(one: Any, two: Any | None = None)"
+
+
+def test_plan_multi_parameter_fallback_help_text(client):
+    plan = Plan(
+        "foo",
+        PlanModel(
+            name="foo",
+            schema={
+                "properties": {
+                    "one": {},
+                    "two": {
+                        "anyOf": [{"items": {}, "type": "array"}, {"type": "boolean"}],
+                    },
+                    "three": {"default": 3},
+                    "four": {"default": None},
+                },
+                "required": ["one", "two"],
+            },
+        ),
+        client,
+    )
+    assert plan.help_text == dedent("""\
+            Plan foo(
+                one: Any,
+                two: list[Any] | bool,
+                three: Any = 3,
+                four: Any | None = None
+            )""")
+
+
+def test_plan_help_text_with_ref(client):
+    schema = {
+        "$defs": {
+            "Spec": {
+                "properties": {
+                    "foo": {"type": "integer"},
+                    "bar": {"$ref": "#/$defs/InnerSpec"},
+                },
+                "required": ["foo", "bar"],
+            },
+            "InnerSpec": {
+                "properties": {
+                    "x": {"type": "number"},
+                    "y": {"default": 10, "type": "number"},
+                },
+                "required": ["x"],
+            },
+        },
+        "properties": {
+            "spec": {"$ref": "#/$defs/Spec"},
+            "meta": {"type": "string", "default": "abc"},
+        },
+        "required": ["spec"],
+    }
+
+    plan = Plan(
+        "ref_plan",
+        PlanModel(name="ref_plan", schema=schema),
+        client,
+    )
+
+    expected = "Plan ref_plan(spec: Spec, meta: str = 'abc')"
+
+    assert plan.help_text == expected
 
 
 def test_plan_properties(client):
@@ -818,8 +883,7 @@ def test_plan_properties(client):
         ),
         client,
     )
-
-    assert plan.properties == {"one", "two"}
+    assert plan.properties == {"one": {}, "two": {}}
     assert plan.required == ["one"]