diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index d927f74ff6..26ed8da781 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from importlib import import_module from inspect import Parameter, signature -from types import ModuleType, UnionType +from types import ModuleType, NoneType, UnionType from typing import Any, Generic, TypeVar, Union, get_args, get_origin, get_type_hints from bluesky.protocols import HasName @@ -12,7 +12,7 @@ from ophyd_async.core import NotConnected from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, create_model from pydantic.fields import FieldInfo -from pydantic.json_schema import JsonSchemaValue +from pydantic.json_schema import JsonSchemaValue, SkipJsonSchema from pydantic_core import CoreSchema, core_schema from blueapi import utils @@ -323,12 +323,12 @@ def _type_spec_for_function( no_default = para.default is Parameter.empty factory = None if no_default else DefaultFactory(para.default) new_args[name] = ( - self._convert_type(arg_type), + self._convert_type(arg_type, no_default), FieldInfo(default_factory=factory), ) return new_args - def _convert_type(self, typ: type | Any) -> type: + def _convert_type(self, typ: type | Any, no_default: bool = True) -> type: """ Recursively convert a type to something that can be deserialised by pydantic. Bluesky protocols (and types that extend them) are replaced @@ -344,12 +344,14 @@ def _convert_type(self, typ: type | Any) -> type: Returns: A Type that can be deserialised by Pydantic """ + if typ is NoneType and not no_default: + return SkipJsonSchema[NoneType] root = get_origin(typ) if is_bluesky_type(typ) or (root is not None and is_bluesky_type(root)): return self._reference(typ) args = get_args(typ) if args: - new_types = tuple(self._convert_type(i) for i in args) + new_types = tuple(self._convert_type(i, no_default) for i in args) if root == UnionType: root = Union return root[new_types] if root else typ # type: ignore diff --git a/tests/system_tests/plans.json b/tests/system_tests/plans.json index bd19637753..4f2075e2b6 100644 --- a/tests/system_tests/plans.json +++ b/tests/system_tests/plans.json @@ -37,16 +37,9 @@ "title": "Delay" }, "metadata": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "title": "Metadata" + "additionalProperties": true, + "title": "Metadata", + "type": "object" } }, "required": [ @@ -89,7 +82,7 @@ }, "radius": { "description": "Radius of the circle", - "exclusiveMinimum": 0, + "exclusiveMinimum": 0.0, "title": "Radius", "type": "number" }, @@ -228,18 +221,18 @@ }, "x_radius": { "description": "The radius along the x axis of the ellipse", - "exclusiveMinimum": 0, + "exclusiveMinimum": 0.0, "title": "X Radius", "type": "number" }, "y_radius": { "description": "The radius along the y axis of the ellipse", - "exclusiveMinimum": 0, + "exclusiveMinimum": 0.0, "title": "Y Radius", "type": "number" }, "angle": { - "default": 0, + "default": 0.0, "description": "The angle of the ellipse (degrees)", "title": "Angle", "type": "number" @@ -510,7 +503,7 @@ "type": "number" }, "angle": { - "default": 0, + "default": 0.0, "description": "Clockwise rotation angle of the rectangle", "title": "Angle", "type": "number" @@ -724,7 +717,7 @@ "type": "integer" }, "rotate": { - "default": 0, + "default": 0.0, "description": "How much to rotate the angle of the spiral", "title": "Rotate", "type": "number" @@ -908,16 +901,9 @@ "$ref": "#/$defs/Spec" }, "metadata": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "title": "Metadata" + "additionalProperties": true, + "title": "Metadata", + "type": "object" } }, "required": [ @@ -946,15 +932,8 @@ "title": "Value" }, "group": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Group" + "title": "Group", + "type": "string" }, "wait": { "title": "Wait", @@ -987,15 +966,8 @@ "title": "Value" }, "group": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Group" + "title": "Group", + "type": "string" }, "wait": { "title": "Wait", @@ -1022,15 +994,8 @@ "type": "object" }, "group": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Group" + "title": "Group", + "type": "string" } }, "required": [ @@ -1052,15 +1017,8 @@ "type": "object" }, "group": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Group" + "title": "Group", + "type": "string" } }, "required": [ @@ -1095,26 +1053,12 @@ "additionalProperties": false, "properties": { "group": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Group" + "title": "Group", + "type": "string" }, "timeout": { - "anyOf": [ - { - "type": "number" - }, - { - "type": "null" - } - ], - "title": "Timeout" + "title": "Timeout", + "type": "number" } }, "title": "wait", diff --git a/tests/unit_tests/core/test_context.py b/tests/unit_tests/core/test_context.py index dab1b03df3..58031178ad 100644 --- a/tests/unit_tests/core/test_context.py +++ b/tests/unit_tests/core/test_context.py @@ -2,6 +2,7 @@ from dataclasses import dataclass from pathlib import Path +from types import NoneType from typing import Generic, TypeVar, Union from unittest.mock import patch @@ -26,6 +27,7 @@ from ophyd_async.epics.adaravis import AravisDetector from ophyd_async.epics.motor import Motor from pydantic import TypeAdapter, ValidationError +from pydantic.json_schema import SkipJsonSchema from pytest import LogCaptureFixture from blueapi.config import EnvironmentConfig, MetadataConfig, Source, SourceKind @@ -431,6 +433,18 @@ def test_reference_type_conversion_new_style_union( assert empty_context._convert_type(Movable | int) == movable_ref | int +def test_reference_type_conversion_new_style_optional( + empty_context: BlueskyContext, +): + movable_ref: type = empty_context._reference(Movable) + assert empty_context._convert_type(Movable) == movable_ref + assert empty_context._convert_type(Movable | None) == movable_ref | None + assert ( + empty_context._convert_type(Movable | None, no_default=False) + == movable_ref | SkipJsonSchema[NoneType] + ) + + def test_default_device_reference(empty_context: BlueskyContext): def default_movable(mov: Movable = inject("demo")) -> MsgGenerator: yield from () @@ -596,3 +610,59 @@ class InnerClass: ... @pytest.mark.parametrize("type,expected", qualified_name_test_data) def test_qualified_name_with_types(type: type, expected: str): assert qualified_name(type) == expected + + +def test_optional_arg_generated_schema( + empty_context: BlueskyContext, +): + def demo_plan(foo: int | None = None) -> MsgGenerator: + yield from () + + empty_context.register_plan(demo_plan) + schema = empty_context.plans["demo_plan"].model.model_json_schema() + assert schema["properties"] == { + "foo": {"title": "Foo", "type": "integer"}, + } + assert "foo" not in schema.get("required", []) + + +def test_overloaded_arg_generated_schema( + empty_context: BlueskyContext, +): + def demo_plan(foo: int | str) -> MsgGenerator: + yield from () + + 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"}]} + } + assert "foo" in schema.get("required", []) + + +def test_optional_overloaded_arg_generated_schema( + empty_context: BlueskyContext, +): + def demo_plan(foo: int | str | None = None) -> MsgGenerator: + yield from () + + 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"}]} + } + assert "foo" not in schema.get("required", []) + + +def test_explicit_none_arg_generated_schema( + empty_context: BlueskyContext, +): + def demo_plan(foo: int | None) -> MsgGenerator: + yield from () + + 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"}]} + } + assert "foo" in schema.get("required", [])