diff --git a/ju/json_schema.py b/ju/json_schema.py index e4a59c6..d3ee115 100644 --- a/ju/json_schema.py +++ b/ju/json_schema.py @@ -23,7 +23,7 @@ """ -from typing import Union, Optional, Any +from typing import Union, Optional, Any, get_args, get_origin from collections import UserString from collections.abc import Mapping, Sequence, Callable from inspect import Parameter @@ -35,7 +35,7 @@ from pydantic import BaseModel from i2 import name_of_obj as _name_of_obj -from ju.util import is_type +from ju.util import is_type, unwrap_optional DFLT_FUNC_TITLE = "" @@ -611,6 +611,54 @@ def _json_safe(value): return True, value +def _array_items(annotation, default, param_to_prop_type: Callable) -> dict | None: + """ + The JSON schema of an array parameter's items, or ``None`` if there is no telling. + + The element type comes from the annotation's arguments (``list[int]``, + ``tuple[int, ...]``, ``Sequence[list[str]]``), else from the default's elements + when they all have the same JSON type. + + >>> from ju.json_schema import DFLT_PARAM_TO_TYPE as to_type + >>> _array_items(list[int], inspect.Parameter.empty, to_type) + {'type': 'integer'} + >>> _array_items(list[list[str]], inspect.Parameter.empty, to_type) + {'type': 'array', 'items': {'type': 'string'}} + >>> _array_items(inspect.Parameter.empty, [1, 2, 3], to_type) + {'type': 'integer'} + >>> _array_items(inspect.Parameter.empty, [1, 'a'], to_type) is None + True + >>> _array_items(list, [], to_type) is None + True + """ + empty = inspect.Parameter.empty + args = get_args(unwrap_optional(annotation)) if annotation is not empty else () + args = tuple(a for a in args if a is not Ellipsis) + if len(args) == 1: # list[int], tuple[int, ...]; not tuple[int, str] + return _value_schema(args[0], empty, param_to_prop_type) + if args: + return None + if isinstance(default, (list, tuple)) and default: + schemas = [_value_schema(empty, d, param_to_prop_type) for d in default] + if all(schema == schemas[0] for schema in schemas): + return schemas[0] + return None + + +def _value_schema(annotation, default, param_to_prop_type: Callable) -> dict: + """``{"type": ..., "items": ...}`` for a value described by an annotation and/or a + default.""" + param = Parameter( + "_", Parameter.KEYWORD_ONLY, annotation=annotation, default=default + ) + schema = {"type": param_to_prop_type(param)} + if schema["type"] == "array": + items = _array_items(annotation, default, param_to_prop_type) + if items is not None: + schema["items"] = items + return schema + + def get_properties(parameters, *, param_to_prop_type: Callable = DFLT_PARAM_TO_TYPE): """ Returns the properties dict for the JSON schema. @@ -644,6 +692,10 @@ def get_properties(parameters, *, param_to_prop_type: Callable = DFLT_PARAM_TO_T name, param = item field = {} field["type"] = param_to_prop_type(param) + if field["type"] == "array": + items = _array_items(param.annotation, param.default, param_to_prop_type) + if items is not None: + field["items"] = items # If there's a default value, add it if param.default is not inspect.Parameter.empty: diff --git a/ju/tests/json_schema_test.py b/ju/tests/json_schema_test.py index b2c0181..3df58d5 100644 --- a/ju/tests/json_schema_test.py +++ b/ju/tests/json_schema_test.py @@ -166,3 +166,77 @@ def test_param_default_to_json_type(default, json_type): param = Parameter("p", Parameter.KEYWORD_ONLY, default=default) assert DFLT_PARAM_TO_TYPE(param) == json_type + + +@pytest.mark.parametrize( + "annotation, default, expected", + [ + (list, Parameter.empty, {"type": "array"}), + (list[int], Parameter.empty, {"type": "array", "items": {"type": "integer"}}), + ( + "tuple[int, ...]", + Parameter.empty, + {"type": "array", "items": {"type": "integer"}}, + ), + ( + list[list[str]], + Parameter.empty, + {"type": "array", "items": {"type": "array", "items": {"type": "string"}}}, + ), + ( + list[float] | None, + Parameter.empty, + {"type": "array", "items": {"type": "number"}}, + ), + ( + Parameter.empty, + [1, 2, 3], + {"type": "array", "items": {"type": "integer"}, "default": [1, 2, 3]}, + ), + ( + Parameter.empty, + ["a", "b"], + {"type": "array", "items": {"type": "string"}, "default": ["a", "b"]}, + ), + # the annotation wins over the default's elements + ( + list[float], + [1, 2], + {"type": "array", "items": {"type": "number"}, "default": [1, 2]}, + ), + # no telling: heterogeneous, empty or fixed-length tuples + (Parameter.empty, [1, "a"], {"type": "array", "default": [1, "a"]}), + (Parameter.empty, [], {"type": "array", "default": []}), + (tuple[int, str], Parameter.empty, {"type": "array"}), + (dict[str, int], Parameter.empty, {"type": "object"}), + (str, Parameter.empty, {"type": "string"}), + ], +) +def test_array_params_from_annotation_and_default(annotation, default, expected): + from ju.json_schema import get_properties + + if annotation == "tuple[int, ...]": + annotation = tuple[int, ...] + param = Parameter( + "p", Parameter.KEYWORD_ONLY, annotation=annotation, default=default + ) + assert get_properties({"p": param}) == {"p": expected} + + +def test_annotated_list_params_in_signature_schema(): + from ju.json_schema import signature_to_json_schema + + def f(xs: list[int], names: list = None, ys=(1.5, 2.5)): ... + + props = signature_to_json_schema(f)["properties"] + assert props["xs"] == {"type": "array", "items": {"type": "integer"}} + assert props["names"]["type"] == "array" and "items" not in props["names"] + assert props["ys"]["items"] == {"type": "number"} + + +@pytest.mark.parametrize("annotation", [bytes, bytearray, memoryview, UserString, str]) +def test_bytes_like_and_str_annotations_stay_string(annotation): + from ju.json_schema import DFLT_PARAM_TO_TYPE + + param = Parameter("p", Parameter.KEYWORD_ONLY, annotation=annotation) + assert DFLT_PARAM_TO_TYPE(param) == "string" diff --git a/ju/util.py b/ju/util.py index d767f49..0305add 100644 --- a/ju/util.py +++ b/ju/util.py @@ -18,7 +18,7 @@ from collections.abc import Mapping, Callable, Sequence from collections import defaultdict from dataclasses import dataclass -from types import GenericAlias +from types import GenericAlias, UnionType from inspect import Parameter from i2 import mk_sentinel @@ -51,6 +51,56 @@ def asis(obj): return obj +def _is_union(tp) -> bool: + return get_origin(tp) in (Union, UnionType) + + +def unwrap_optional(annotation): + """ + Strips ``None`` out of an ``Optional``/union annotation; returns the annotation + unchanged if it is not a union, and the single remaining member if there is one. + + >>> from typing import Optional + >>> unwrap_optional(Optional[list[int]]) + list[int] + >>> unwrap_optional(int | None) + + >>> unwrap_optional(int | str | None) + typing.Union[int, str] + >>> unwrap_optional(list[int]) + list[int] + """ + if not _is_union(annotation): + return annotation + members = tuple(a for a in get_args(annotation) if a is not type(None)) + if len(members) == 1: + return members[0] + return Union[members] if members else annotation + + +def _annotation_matches(annotation, type_) -> bool: + """Does ``annotation`` (a class or a subscripted generic) denote ``type_``, or a + subclass of it? If ``type_`` is subscripted, the annotation must have the same + arguments. (The order of a type mapping matters: a parent listed before its child + claims the child's annotations too.)""" + annotation = unwrap_optional(annotation) + if _is_union(annotation): + return False + origin = get_origin(annotation) or annotation + type_origin = get_origin(type_) or type_ + if not (isinstance(origin, type) and isinstance(type_origin, type)): + return False + try: + if not issubclass(origin, type_origin): + return False + except TypeError: # e.g. a Protocol with data members + return False + type_args, args = get_args(type_), get_args(annotation) + if type_args: # ``bytes`` is a ``Sequence``, but not a ``Sequence[str]`` + return args == type_args + return True + + def is_type(param: Parameter, type_: SomeType): """ Checks if the type of a parameter's default value or its annotation matches a @@ -95,9 +145,31 @@ def is_type(param: Parameter, type_: SomeType): True >>> is_type(param, Sequence[str]) False + + An annotation alone is enough, bare or subscripted, optional or not -- it matches + ``type_`` or any subclass of it: + + >>> annotated = Parameter('p', Parameter.KEYWORD_ONLY, annotation=list[int]) + >>> is_type(annotated, Sequence) + True + >>> is_type(annotated, list) + True + >>> is_type(annotated, List[int]) + True + >>> is_type(annotated, List[str]) + False + >>> is_type(annotated, str) + False + >>> optional = Parameter('p', Parameter.KEYWORD_ONLY, annotation=list | None) + >>> is_type(optional, Sequence) + True """ if param.annotation is type_: return True + if param.annotation is not Parameter.empty and _annotation_matches( + param.annotation, type_ + ): + return True if isinstance(type_, type) and not hasattr(type_, "__origin__"): return isinstance(param.default, type_) if hasattr(type_, "__origin__"):