diff --git a/src/blueapi/client/client.py b/src/blueapi/client/client.py index 7e95154466..66a79d8dd6 100644 --- a/src/blueapi/client/client.py +++ b/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,29 @@ 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" + + args = [_format_arg(name, info) for name, info in self.properties.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: @@ -806,3 +834,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" diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 2e5d67de24..76ce74e647 100644 --- a/src/blueapi/core/context.py +++ b/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 diff --git a/tests/system_tests/plans.json b/tests/system_tests/plans.json index 3a3bfbe3fb..6f84a9c389 100644 --- a/tests/system_tests/plans.json +++ b/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" } diff --git a/tests/unit_tests/client/test_client.py b/tests/unit_tests/client/test_client.py index cbf41c87b0..d5b0493adc 100644 --- a/tests/unit_tests/client/test_client.py +++ b/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"] diff --git a/tests/unit_tests/core/test_context.py b/tests/unit_tests/core/test_context.py index 6a39c68f5e..d8da5c16f7 100644 --- a/tests/unit_tests/core/test_context.py +++ b/tests/unit_tests/core/test_context.py @@ -1,9 +1,10 @@ from __future__ import annotations from dataclasses import dataclass, field +from inspect import Parameter from pathlib import Path from types import ModuleType, NoneType -from typing import Any, Generic, TypeVar, Union +from typing import Any, Generic, TypeVar, Union, get_args, get_type_hints from unittest.mock import MagicMock, Mock, patch import pytest @@ -45,7 +46,7 @@ TiledConfig, ) from blueapi.core import BlueskyContext, is_bluesky_compatible_device -from blueapi.core.context import DefaultFactory, generic_bounds, qualified_name +from blueapi.core.context import generic_bounds, qualified_name from blueapi.core.protocols import DeviceConnectResult, DeviceManager from blueapi.utils.invalid_config_error import InvalidConfigError @@ -86,6 +87,10 @@ def has_some_params(foo: int = 42, bar: str = "bar") -> MsgGenerator: yield from () +def has_optional_parameter(foo: dict[str, Any] | None = None) -> MsgGenerator: + yield from () + + def has_typeless_param(foo) -> MsgGenerator: yield from () @@ -167,7 +172,9 @@ def some_configurable() -> SomeConfigurable: return SomeConfigurable() -@pytest.mark.parametrize("plan", [has_no_params, has_one_param, has_some_params]) +@pytest.mark.parametrize( + "plan", [has_no_params, has_one_param, has_some_params, has_optional_parameter] +) def test_add_plan(empty_context: BlueskyContext, plan: PlanGenerator): empty_context.register_plan(plan) assert plan.__name__ in empty_context.plans @@ -354,12 +361,23 @@ def test_add_metadata_with_config( assert md in empty_context.run_engine.md.items() -def test_function_spec(empty_context: BlueskyContext): +def test_function_spec_with_some_params(empty_context: BlueskyContext): spec = empty_context._type_spec_for_function(has_some_params) assert spec["foo"][0] is int - assert spec["foo"][1].default_factory == DefaultFactory(42) + assert spec["foo"][1].default == 42 assert spec["bar"][0] is str - assert spec["bar"][1].default_factory == DefaultFactory("bar") + assert spec["bar"][1].default == "bar" + + +def test_function_spec_with_optional_params(empty_context: BlueskyContext): + spec = empty_context._type_spec_for_function(has_optional_parameter) + types = get_type_hints(has_optional_parameter) + arg_type = types.get("foo", Parameter.empty) + + _type = SkipJsonSchema[empty_context._convert_type(arg_type, False)] + inner_type, *annotations = get_args(_type) + assert spec["foo"][0] == inner_type + assert spec["foo"][1].default_factory() is None # type: ignore def test_basic_type_conversion(empty_context: BlueskyContext): @@ -440,7 +458,7 @@ def default_movable(mov: Movable = inject("demo")) -> MsgGenerator: spec = empty_context._type_spec_for_function(default_movable) movable_ref = empty_context._reference(Movable) assert spec["mov"][0] == movable_ref - assert spec["mov"][1].default_factory == DefaultFactory("demo") + assert spec["mov"][1].default == "demo" def test_generic_default_device_reference(empty_context: BlueskyContext): @@ -450,7 +468,7 @@ def default_movable(mov: Movable[float] = inject("demo")) -> MsgGenerator: spec = empty_context._type_spec_for_function(default_movable) motor_ref = empty_context._reference(Movable[float]) assert spec["mov"][0] == motor_ref - assert spec["mov"][1].default_factory == DefaultFactory("demo") + assert spec["mov"][1].default == "demo" class ConcreteStoppable(Stoppable): @@ -500,7 +518,7 @@ def test_str_default(empty_context: BlueskyContext, sim_motor: Motor, alt_motor: spec = empty_context._type_spec_for_function(has_default_reference) assert spec["m"][0] is movable_ref - assert (df := spec["m"][1].default_factory) and df() == SIM_MOTOR_NAME # type: ignore + assert spec["m"][1].default == SIM_MOTOR_NAME assert has_default_reference.__name__ in empty_context.plans model = empty_context.plans[has_default_reference.__name__].model @@ -519,7 +537,7 @@ def test_nested_str_default( spec = empty_context._type_spec_for_function(has_default_nested_reference) assert spec["m"][0] == list[movable_ref] - assert (df := spec["m"][1].default_factory) and df() == [SIM_MOTOR_NAME] # type: ignore + assert spec["m"][1].default == [SIM_MOTOR_NAME] assert has_default_nested_reference.__name__ in empty_context.plans model = empty_context.plans[has_default_nested_reference.__name__].model @@ -665,7 +683,10 @@ def demo_plan(foo: int | None) -> MsgGenerator: empty_context.register_plan(demo_plan) schema = empty_context.plans["demo_plan"].model.model_json_schema() assert schema["properties"] == { - "foo": {"title": "Foo", "anyOf": [{"type": "integer"}, {"type": "null"}]} + "foo": { + "title": "Foo", + "anyOf": [{"type": "integer"}, {"type": "null"}], + } } assert "foo" in schema.get("required", [])