diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index d4cf0e809..0c76542fa 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -3,7 +3,7 @@ from inspect import Parameter, signature from typing import Any, Self, TypedDict -from pydantic import Field, TypeAdapter, ValidationError +from pydantic import Field, ValidationError from blueapi.core import BlueskyContext from blueapi.utils import BlueapiBaseModel @@ -61,43 +61,53 @@ def prepare_params( def _inner_prepare_params( self, ctx: BlueskyContext ) -> tuple[list[Any], Mapping[str, Any]]: + """Prepare positional and keyword arguments for a plan invocation. + + Bind the task's supplied arguments to the plan's signature, then validate + them using the plan's Pydantic model. This also resolves injected device + names to the corresponding devices in the context and applies parameter + defaults defined by the model. + + Construct the positional and keyword arguments according to the plan's + parameter kinds, ensuring that positional-only, variadic, and + keyword-only parameters are handled correctly. + + Args: + ctx: The context containing the registered plan and its associated + parameter model. + + Returns: + A tuple containing the prepared positional and keyword arguments + to pass to the plan. + """ plan = ctx.plans[self.name] func = ctx.plan_functions[self.name] sig = signature(func) bound = sig.bind_partial(*self.args, **self.params) - # Only validate explicitly supplied arguments. This allows Pydantic's - # default/default_factory to provide injected defaults. - adapter = TypeAdapter(plan.model) - validated = adapter.validate_python(bound.arguments) + validated = plan.model.model_validate(bound.arguments) args: list[Any] = [] kwargs: dict[str, Any] = {} for name, parameter in sig.parameters.items(): - supplied = name in bound.arguments - - if supplied: - value = getattr(validated, name) - - match parameter.kind: - case Parameter.POSITIONAL_ONLY: + value = getattr(validated, name) + + match parameter.kind: + case Parameter.POSITIONAL_ONLY: + args.append(value) + case Parameter.VAR_POSITIONAL: + args.extend(value) + case Parameter.VAR_KEYWORD: + kwargs.update(value) + case Parameter.POSITIONAL_OR_KEYWORD: + if name in bound.arguments and name not in self.params: args.append(value) - case Parameter.VAR_POSITIONAL: - args.extend(value) - case Parameter.POSITIONAL_OR_KEYWORD: - if name in self.params: - kwargs[name] = value - else: - args.append(value) - case Parameter.KEYWORD_ONLY: + else: kwargs[name] = value - case Parameter.VAR_KEYWORD: - kwargs.update(value) - else: - # Let the generated Pydantic model provide the default. - kwargs[name] = getattr(validated, name) + case Parameter.KEYWORD_ONLY: + kwargs[name] = value return args, kwargs diff --git a/tests/unit_tests/worker/test_task_worker.py b/tests/unit_tests/worker/test_task_worker.py index 8f0a1269c..044668d93 100644 --- a/tests/unit_tests/worker/test_task_worker.py +++ b/tests/unit_tests/worker/test_task_worker.py @@ -1115,7 +1115,6 @@ def positional_only_plan(val: int, /, **devices: FakeDevice) -> MsgGenerator: def test_arg_in_args_and_kwargs(): - def positionals(one: int, two: int) -> MsgGenerator: yield from [] @@ -1127,3 +1126,74 @@ def positionals(one: int, two: int) -> MsgGenerator: task.prepare_params(context) assert tpe.value.errors[0]["msg"] == "multiple values for argument 'one'" + + +@pytest.mark.parametrize( + "extra_args", + [[], [2], [2, 3]], + ids=["no_extra_args", "one_extra_arg", "two_extra_args"], +) +@pytest.mark.parametrize( + "keyword_only_params", + [{}, {"group": "test"}, {"wait": True}, {"group": "test", "wait": True}], + ids=["no_keyword_only_params", "with_group", "with_wait", "with_group_and_wait"], +) +@pytest.mark.parametrize( + "variadic_kwargs", + [{}, {"custom": "value"}, {"custom": "value", "another": 42}], + ids=["no_variadic_kwargs", "one_variadic_kwarg", "two_variadic_kwargs"], +) +def test_plan_args_with_keyword_only_params_and_variadic_kwargs( + fake_device: FakeDevice, + extra_args: list[int], + keyword_only_params: dict[str, Any], + variadic_kwargs: dict[str, Any], +): + def args_with_keyword_only_params_and_variadic_kwargs( + obj: FakeDevice, + *args: Any, + group: str | None = None, + wait: bool = False, + **kwargs: Any, + ) -> MsgGenerator: + yield from [] + + context = BlueskyContext() + context.register_plan(args_with_keyword_only_params_and_variadic_kwargs) + context.register_device(fake_device) + + task = Task( + name="args_with_keyword_only_params_and_variadic_kwargs", + args=[fake_device.name, 1, *extra_args], + params={**keyword_only_params, **variadic_kwargs}, + ) + prepared_args, prepared_kwargs = task.prepare_params(context) + + expected_args = [fake_device, 1, *extra_args] + expected_kwargs = { + "group": None, # Defaults supplied by schema from registering plan. + "wait": False, # Defaults supplied by schema from registering plan. + **keyword_only_params, + **variadic_kwargs, + } + assert prepared_args == expected_args + assert prepared_kwargs == expected_kwargs + + +def test_string_default_is_not_treated_as_device_injection(): + def plan_with_default_str(arg: str = "primary") -> MsgGenerator: + yield from [] + + primary = FakeDevice("primary") + context = BlueskyContext() + context.register_device(primary) + context.register_plan(plan_with_default_str) + + task = Task(name="plan_with_default_str") + + args, kwargs = task.prepare_params(context) + assert args == [] + # Defaults provided by plan. Checking that device with same name is not added + # by mistake as the default plan value. The kwargs is provided by pydantic model + # from schema. + assert kwargs == {"arg": "primary"}