Infer array 'items' and recognise annotation-only list params (#19)

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:

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"

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)
+    <class 'int'>
+    >>> 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__"):