diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index c5b41ff45..d62b97e1a 100644 --- a/src/blueapi/client/client.py +++ b/src/blueapi/client/client.py @@ -4,9 +4,8 @@ from collections.abc import Iterable from concurrent.futures import Future from functools import cached_property -from itertools import chain 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 @@ -50,6 +49,32 @@ log = logging.getLogger(__name__) +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") + type_map = { + "string": "str", + "integer": "int", + "boolean": "bool", + "number": "float", + "object": "dict", + } + if isinstance(json_type, str): + return type_map.get(json_type, json_type.split(".")[-1]) + + return "Any" + + class MissingInstrumentSessionError(Exception): pass @@ -154,7 +179,7 @@ def help_text(self) -> str: return self.model.description or f"Plan {self!r}" @property - def properties(self) -> set[str]: + def properties(self) -> dict[str, Any]: return self.model.parameter_schema.get("properties", {}).keys() @property @@ -192,9 +217,26 @@ def _build_args(self, *args, **kwargs): 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})" + props = self.model.parameter_schema.get("properties", {}) + tab = " " + args = [] + for name, info in props.items(): + typ = _pretty_type(info) + arg = f"{name}: {typ}" + if name not in self.required: + arg = f"{arg} | None = None" + args.append(arg) + + single_line = f"{self.name}({', '.join(args)})" + max_length = 100 + max_args_inline = 3 + if len(single_line) <= max_length and len(args) <= max_args_inline: + return single_line + + # Fall back to multiline if too many arguments or too long. + multiline_args = ",\n".join(f"{tab}{arg}" for arg in args) + + return f"{self.name}(\n{multiline_args}\n)" class BlueapiClient: diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index a96f428e8..f79b8aaef 100644 --- a/tests/unit_tests/client/test_client.py +++ b/tests/unit_tests/client/test_client.py @@ -716,7 +716,36 @@ 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": {}, + "four": {}, + }, + "required": ["one", "two"], + }, + ), + client, + ) + assert ( + plan.help_text == "Plan foo(\n" + " one: Any,\n" + " two: list[Any] | bool,\n" + " three: Any | None = None,\n" + " four: Any | None = None\n" + ")" + ) def test_plan_properties(client):