Schema inference: resolve string annotations, Annotated and NewType (#27)

ju/json_schema.py

@@ -35,7 +35,7 @@
 from pydantic import BaseModel
 from i2 import name_of_obj as _name_of_obj
 
-from ju.util import is_type, unwrap_optional
+from ju.util import is_type, unwrap_optional, resolve_annotations
 
 DFLT_FUNC_TITLE = ""
 
@@ -316,6 +316,8 @@ def signature_to_json_schema(
         func_name = name_of_obj(func)
 
     parameters = sig.parameters
+    if callable(func):  # (``from __future__ import annotations`` -> strings)
+        parameters = resolve_annotations(func, parameters)
 
     schema = deepcopy(BASE_SCHEMA)
     schema["title"] = pyname_to_title(func_name)
@@ -611,7 +613,17 @@ def _json_safe(value):
     return True, value
 
 
-def _array_items(annotation, default, param_to_prop_type: Callable) -> dict | None:
+_MAX_ITEMS_DEPTH = 16  # (a recursive ``type`` alias would nest ``items`` forever)
+
+
+def _array_items(
+    annotation,
+    default,
+    param_to_prop_type: Callable,
+    name: str = "_",
+    *,
+    depth: int = 0,
+) -> dict | None:
     """
     The JSON schema of an array parameter's items, or ``None`` if there is no telling.
 
@@ -632,28 +644,43 @@ def _array_items(annotation, default, param_to_prop_type: Callable) -> dict | No
     True
     """
     empty = inspect.Parameter.empty
+    if depth >= _MAX_ITEMS_DEPTH:
+        return None
     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)
+        return _value_schema(args[0], empty, param_to_prop_type, name, depth=depth + 1)
     if args:
         return None
     if isinstance(default, (list, tuple)) and default:
-        schemas = [_value_schema(empty, d, param_to_prop_type) for d in default]
+        schemas = [
+            _value_schema(empty, d, param_to_prop_type, name, depth=depth + 1)
+            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:
+def _value_schema(
+    annotation,
+    default,
+    param_to_prop_type: Callable,
+    name: str = "_",
+    *,
+    depth: int = 0,
+) -> dict:
     """``{"type": ..., "items": ...}`` for a value described by an annotation and/or a
-    default."""
+    default. ``name`` is that of the parameter the value belongs to (so a custom
+    ``param_to_prop_type`` sees it, not a placeholder)."""
     param = Parameter(
-        "_", Parameter.KEYWORD_ONLY, annotation=annotation, default=default
+        name, 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)
+        items = _array_items(
+            annotation, default, param_to_prop_type, name, depth=depth
+        )
         if items is not None:
             schema["items"] = items
     return schema
@@ -693,7 +720,9 @@ def get_properties(parameters, *, param_to_prop_type: Callable = DFLT_PARAM_TO_T
         field = {}
         field["type"] = param_to_prop_type(param)
         if field["type"] == "array":
-            items = _array_items(param.annotation, param.default, param_to_prop_type)
+            items = _array_items(
+                param.annotation, param.default, param_to_prop_type, name
+            )
             if items is not None:
                 field["items"] = items
 

ju/tests/annotations_test.py

@@ -0,0 +1,115 @@
+"""Tests for schema inference from string, ``Annotated`` and ``NewType`` annotations
+(issue i2mint/ju#27). This module is itself under ``from __future__ import
+annotations``, so every annotation in it is a string."""
+
+from __future__ import annotations
+
+import sys
+from functools import partial
+from inspect import Parameter
+from typing import Annotated, NewType, Optional, TYPE_CHECKING
+
+import pytest
+
+from ju.json_schema import DFLT_PARAM_TO_TYPE, signature_to_json_schema
+
+if TYPE_CHECKING:
+    from nowhere import Unimportable
+
+UserId = NewType("UserId", int)
+Tags = Annotated[list[str], "some metadata"]
+
+
+def _props(func):
+    return signature_to_json_schema(func)["properties"]
+
+
+def test_string_annotations_are_resolved():
+    def f(n: int, xs: list[int], flag: bool = True, opt: Optional[float] = None): ...
+
+    assert _props(f) == {
+        "n": {"type": "integer"},
+        "xs": {"type": "array", "items": {"type": "integer"}},
+        "flag": {"type": "boolean", "default": True},
+        "opt": {"type": "number", "default": None},
+    }
+
+
+def test_annotated_and_newtype():
+    def f(a: Annotated[list[int], "m"], b: UserId, c: Tags, d: Annotated[UserId, 1]): ...
+
+    assert _props(f) == {
+        "a": {"type": "array", "items": {"type": "integer"}},
+        "b": {"type": "integer"},
+        "c": {"type": "array", "items": {"type": "string"}},
+        "d": {"type": "integer"},
+    }
+
+
+def test_wrappers_without_string_annotations():
+    # (a real, non-string annotation: built from the objects themselves)
+    param = Parameter(
+        "p", Parameter.KEYWORD_ONLY, annotation=Optional[Annotated[UserId, "m"]]
+    )
+    assert DFLT_PARAM_TO_TYPE(param) == "integer"
+
+
+def test_unresolvable_annotation_falls_back_without_spoiling_the_others():
+    def f(a: int, b: Unimportable, c: float): ...
+
+    assert _props(f) == {
+        "a": {"type": "integer"},
+        "b": {"type": "string"},
+        "c": {"type": "number"},
+    }
+
+
+def test_partial_and_callable_instance_and_method():
+    def g(a: int, b: list[float]): ...
+
+    class Callme:
+        def __call__(self, x: float, ys: list[int]): ...
+
+        def meth(self, z: bool): ...
+
+    assert _props(partial(g, 1)) == {"b": {"type": "array", "items": {"type": "number"}}}
+    assert _props(Callme()) == {
+        "x": {"type": "number"},
+        "ys": {"type": "array", "items": {"type": "integer"}},
+    }
+    assert _props(Callme().meth) == {"z": {"type": "boolean"}}
+
+
+def test_custom_param_to_prop_type_sees_the_owning_name():
+    seen = []
+
+    def spy(param):
+        seen.append(param.name)
+        return DFLT_PARAM_TO_TYPE(param)
+
+    def f(xs: list[list[int]], ys=(1, 2)): ...
+
+    signature_to_json_schema(f, param_to_prop_type=spy)
+    assert set(seen) == {"xs", "ys"}
+
+
+@pytest.mark.skipif(sys.version_info < (3, 12), reason="`type` statement")
+def test_recursive_and_cyclic_type_aliases_terminate():
+    ns = {}
+    exec(
+        "type Tree = list[Tree]\ntype A = B\ntype B = A\n"
+        "def t(x: Tree): ...\ndef c(y: A): ...\n",
+        ns,
+    )
+    assert _props(ns["t"])["x"]["type"] == "array"  # (items nest to a bounded depth)
+    assert _props(ns["c"]) == {"y": {"type": "string"}}
+
+
+def test_class_with_only_a_new_is_resolved():
+    class N:
+        def __new__(cls, a: int, b: list[float]): ...
+
+    assert _props(N) == {
+        "a": {"type": "integer"},
+        "b": {"type": "array", "items": {"type": "number"}},
+    }

ju/util.py

@@ -19,6 +19,9 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from types import GenericAlias, UnionType
+import inspect
+import builtins
+import typing
 from inspect import Parameter
 
 from i2 import mk_sentinel
@@ -55,12 +58,41 @@ def _is_union(tp) -> bool:
     return get_origin(tp) in (Union, UnionType)
 
 
+def _strip_wrappers(annotation):
+    """
+    Peels the wrappers that don't change what a value *is* off an annotation:
+    ``Annotated[T, ...]`` (to ``T``), ``NewType`` (to its supertype) and a ``type``
+    statement alias (to its value). Unions are handled by `unwrap_optional`.
+
+    >>> from typing import Annotated, NewType
+    >>> _strip_wrappers(Annotated[list[int], "meta"])
+    list[int]
+    >>> UserId = NewType("UserId", int)
+    >>> _strip_wrappers(Annotated[UserId, "meta"])
+    <class 'int'>
+    """
+    seen = []  # (aliases can refer to each other: don't loop)
+    while True:
+        if any(annotation is a for a in seen):
+            return annotation
+        seen.append(annotation)
+        if get_origin(annotation) is typing.Annotated:
+            annotation = get_args(annotation)[0]
+        elif hasattr(annotation, "__supertype__"):  # NewType
+            annotation = annotation.__supertype__
+        elif isinstance(annotation, getattr(typing, "TypeAliasType", ())):
+            annotation = annotation.__value__
+        else:
+            return annotation
+
+
 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.
+    ``Annotated`` and ``NewType`` wrappers are peeled off too (also on the members).
 
-    >>> from typing import Optional
+    >>> from typing import Optional, Annotated, NewType
     >>> unwrap_optional(Optional[list[int]])
     list[int]
     >>> unwrap_optional(int | None)
@@ -69,15 +101,74 @@ def unwrap_optional(annotation):
     typing.Union[int, str]
     >>> unwrap_optional(list[int])
     list[int]
+    >>> unwrap_optional(Annotated[Optional[list[int]], "meta"])
+    list[int]
+    >>> unwrap_optional(Optional[NewType("UserId", int)])
+    <class 'int'>
     """
+    annotation = _strip_wrappers(annotation)
     if not _is_union(annotation):
         return annotation
-    members = tuple(a for a in get_args(annotation) if a is not type(None))
+    members = tuple(
+        _strip_wrappers(a) for a in get_args(annotation) if a is not type(None)
+    )
     if len(members) == 1:
-        return members[0]
+        return unwrap_optional(members[0])
     return Union[members] if members else annotation
 
 
+def _annotation_target(func):
+    """The function whose globals/``__annotations__`` describe ``func``'s signature."""
+    while isinstance(func, partial):
+        func = func.func
+    if inspect.isclass(func):
+        if func.__init__ is object.__init__ and func.__new__ is not object.__new__:
+            return func.__new__  # (only a ``__new__`` describes its signature)
+        return func.__init__
+    if not (inspect.isroutine(func) or hasattr(func, "__wrapped__")):
+        return getattr(type(func), "__call__", func)  # a callable instance
+    return func
+
+
+def resolve_annotations(func, parameters):
+    """
+    Returns ``parameters`` (a mapping or iterable of `Parameter`) with the string
+    annotations of ``func`` evaluated -- all of them under
+    ``from __future__ import annotations``. An annotation that can't be resolved
+    (e.g. a name only imported under ``TYPE_CHECKING``) is left as it was.
+
+    >>> def f(a: "int", b: "list[float]" = None, c: "Nope" = 1): ...
+    >>> sig = inspect.signature(f)
+    >>> [(n, p.annotation) for n, p in resolve_annotations(f, sig.parameters).items()]
+    [('a', <class 'int'>), ('b', list[float]), ('c', 'Nope')]
+    """
+    as_mapping = isinstance(parameters, Mapping)
+    params = dict(parameters) if as_mapping else {p.name: p for p in parameters}
+    if not any(isinstance(p.annotation, str) for p in params.values()):
+        return parameters
+    target = _annotation_target(func)
+    try:
+        hints = typing.get_type_hints(target, include_extras=True)
+    except Exception:  # one unresolvable annotation spoils the lot: go one by one
+        hints = {}
+        globalns = {**vars(builtins), **getattr(inspect.unwrap(target), "__globals__", {})}
+        for name, p in params.items():
+            if isinstance(p.annotation, str):
+                try:
+                    hints[name] = eval(p.annotation, globalns)
+                except Exception:
+                    pass
+    resolved = {
+        name: (
+            p.replace(annotation=hints[name])
+            if isinstance(p.annotation, str) and name in hints
+            else p
+        )
+        for name, p in params.items()
+    }
+    return resolved if as_mapping else list(resolved.values())
+
+
 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