Improve type hints for plan signature to include defaults

Author: oliwenmandiamond

src/blueapi/client/client.py

@@ -218,24 +218,33 @@ def _build_args(self, *args, **kwargs):
 
     def __repr__(self):
         props = self.model.parameter_schema.get("properties", {})
+        required = set(self.required)
+
         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"
+
+            if name not in required:
+                if "default" in info:
+                    default = repr(info["default"])
+                    arg = f"{arg} = {default}"
+                else:
+                    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)"
 
 

src/blueapi/core/context.py

@@ -5,7 +5,7 @@
 from importlib import import_module
 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
@@ -516,14 +516,16 @@ def _type_spec_for_function(
             ):
                 default_factory = self._composite_factory(arg_type)
                 _type = SkipJsonSchema[self._convert_type(arg_type, no_default)]
+                field_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),
-            )
+                if no_default:
+                    field_info = FieldInfo()
+                else:
+                    field_info = FieldInfo(default=para.default)
+
+            new_args[name] = (_type, field_info)
+
         return new_args
 
     def _convert_type(self, typ: Any, no_default: bool = True) -> type:
@@ -574,19 +576,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/unit_tests/client/test_client.py

@@ -730,8 +730,8 @@ def test_plan_multi_parameter_fallback_help_text(client):
                     "two": {
                         "anyOf": [{"items": {}, "type": "array"}, {"type": "boolean"}],
                     },
-                    "three": {},
-                    "four": {},
+                    "three": {"default": 3},
+                    "four": {"default": None},
                 },
                 "required": ["one", "two"],
             },
@@ -742,8 +742,8 @@ def test_plan_multi_parameter_fallback_help_text(client):
         plan.help_text == "Plan foo(\n"
         "    one: Any,\n"
         "    two: list[Any] | bool,\n"
-        "    three: Any | None = None,\n"
-        "    four: Any | None = None\n"
+        "    three: Any = 3,\n"
+        "    four: Any = None\n"
         ")"
     )
 

tests/unit_tests/core/test_context.py

@@ -46,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.connect_devices import _establish_device_connections
 from blueapi.utils.invalid_config_error import InvalidConfigError
@@ -431,9 +431,9 @@ def test_with_config_passes_mock_to_with_dodal_module(
 def test_function_spec(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_basic_type_conversion(empty_context: BlueskyContext):
@@ -514,7 +514,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):
@@ -524,7 +524,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):
@@ -574,7 +574,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
@@ -593,7 +593,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
@@ -697,7 +697,7 @@ def demo_plan(foo: int | None = None) -> MsgGenerator:
     empty_context.register_plan(demo_plan)
     schema = empty_context.plans["demo_plan"].model.model_json_schema()
     assert schema["properties"] == {
-        "foo": {"title": "Foo", "type": "integer"},
+        "foo": {"title": "Foo", "type": "integer", "default": None},
     }
     assert "foo" not in schema.get("required", [])
 
@@ -725,7 +725,11 @@ def demo_plan(foo: int | str | None = 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": "string"}]}
+        "foo": {
+            "title": "Foo",
+            "anyOf": [{"type": "integer"}, {"type": "string"}],
+            "default": None,
+        }
     }
     assert "foo" not in schema.get("required", [])
 
@@ -739,7 +743,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", [])