Repository: jcrist/msgspec (python)

Issue:
Support json-schema generation
It would be helpful if JSON schema generation was supported

```python
class User(msgspec.Struct):
    """A new type describing a User"""
    name: str
    groups: Set[str] = set()
    email: Optional[str] = None

schema = User.json_schema()
```

Similar to the functionality seen in https://github.com/s-knibbs/dataclasses-jsonschema

Pull request description:
Generate JSON Schemas from types
This adds two functions:

- ``msgspec.json.schema`` for generating a single JSON Schema from a
  given type.
- ``msgspec.json.schema_components`` for generating multiple schemas
  (and a corresponding components mapping) from multiple types.

~Still needs docs and tests.~

Fixes #125.

Patch:
diff --git a/docs/source/api.rst b/docs/source/api.rst
index 405f7e7..274588b 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -62,6 +62,16 @@ MessagePack
 .. autofunction:: decode
 
 
+JSON Schema
+-----------
+
+.. currentmodule:: msgspec.json
+
+.. autofunction:: schema
+
+.. autofunction:: schema_components
+
+
 Exceptions
 ----------
 
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 87d5944..ee097b7 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -159,6 +159,7 @@ is required.
     structs.rst
     constraints.rst
     extending.rst
+    jsonschema.rst
     perf-tips.rst
     benchmarks.rst
     examples/index.rst
diff --git a/docs/source/jsonschema.rst b/docs/source/jsonschema.rst
new file mode 100644
index 0000000..10504b6
--- /dev/null
+++ b/docs/source/jsonschema.rst
@@ -0,0 +1,98 @@
+JSON Schema
+===========
+
+``msgspec`` provides a few utilities for generating `JSON Schema`_
+specifications from msgspec-compatible :ref:`types <supported-types>` and
+:doc:`constraints <constraints>`.
+
+- `msgspec.json.schema`: generates a complete JSON Schema for a single type.
+- `msgspec.json.schema_components`: generates JSON schemas for multiple types,
+  along with a corresponding ``components`` mapping. This is mainly useful when
+  generating multiple schemas to include in a larger specification like OpenAPI_.
+
+
+The generated schemas are compatible with `JSON Schema`_ 2020-12 and OpenAPI_
+3.1.
+
+
+Example
+-------
+
+
+.. code-block:: python
+
+    import msgspec
+    from msgspec import Struct, Meta
+    from typing import Annotated, Optional
+
+
+    # A float constrained to values > 0
+    PositiveFloat = Annotated[float, Meta(gt=0)]
+
+
+    class Dimensions(Struct):
+        """Dimensions for a product, all measurements in centimeters"""
+        length: PositiveFloat
+        width: PositiveFloat
+        height: PositiveFloat
+
+
+    class Product(Struct):
+        """A product in a catalog"""
+        id: int
+        name: str
+        price: PositiveFloat
+        tags: set[str] = set()
+        dimensions: Optional[Dimensions] = None
+
+
+    # Generate a schema for a list of products
+    schema = msgspec.json.schema(list[Product])
+
+    # Print out that schema as JSON
+    print(msgspec.json.encode(schema))
+
+
+.. code-block:: json
+
+    {
+      "type": "array",
+      "items": {"$ref": "#/$defs/Product"},
+      "$defs": {
+        "Dimensions": {
+          "title": "Dimensions",
+          "description": "Dimensions for a product, all measurements in centimeters",
+          "type": "object",
+          "properties": {
+            "length": {"type": "number", "exclusiveMinimum": 0},
+            "width": {"type": "number", "exclusiveMinimum": 0},
+            "height": {"type": "number", "exclusiveMinimum": 0}
+          },
+          "required": ["length", "width", "height"]
+        },
+        "Product": {
+          "title": "Product",
+          "description": "A product in a catalog",
+          "type": "object",
+          "properties": {
+            "id": {"type": "integer"},
+            "name": {"type": "string"},
+            "price": {"type": "number", "exclusiveMinimum": 0},
+            "tags": {
+              "type": "array",
+              "items": {"type": "string"},
+              "default": [],
+            },
+            "dimensions": {
+              "anyOf": [{"type": "null"}, {"$ref": "#/$defs/Dimensions"}],
+              "default": null,
+            }
+          },
+          "required": ["id", "name", "price"]
+        }
+      }
+    }
+
+
+.. _JSON Schema: https://json-schema.org/
+.. _OpenAPI: https://www.openapis.org/
diff --git a/msgspec/_core.c b/msgspec/_core.c
index 4e00e86..4dad5ca 100644
--- a/msgspec/_core.c
+++ b/msgspec/_core.c
@@ -3860,6 +3860,19 @@ typenode_collect_type_full(
                 state, MS_TYPE_VARTUPLE, PyTuple_GET_ITEM(args, 0)
             );
         }
+        else if (
+            PyTuple_GET_SIZE(args) == 1 &&
+            PyTuple_CheckExact(PyTuple_GET_ITEM(args, 0)) &&
+            PyTuple_GET_SIZE(PyTuple_GET_ITEM(args, 0)) == 0
+        ) {
+            /* XXX: this case handles a weird compatibility issue:
+             * - Tuple[()].__args__ == ((),)
+             * - tuple[()].__args__ == ()
+             */
+            out = typenode_collect_array(
+                state, MS_TYPE_FIXTUPLE, PyTuple_GET_ITEM(args, 0)
+            );
+        }
         else {
             out = typenode_collect_array(state, MS_TYPE_FIXTUPLE, args);
         }
diff --git a/msgspec/_utils.py b/msgspec/_utils.py
index 65e4d59..0ffe451 100644
--- a/msgspec/_utils.py
+++ b/msgspec/_utils.py
@@ -1,3 +1,19 @@
+import datetime
+import enum
+import re
+import sys
+from typing import Any, Union, Literal, List, Tuple, Dict, Set, FrozenSet
+
+if sys.version_info >= (3, 9):
+    from collections.abc import Iterable
+else:
+    from typing import Iterable
+
+try:
+    from types import UnionType
+except Exception:
+    UnionType = None
+
 try:
     from typing_extensions import _AnnotatedAlias
 except Exception:
@@ -21,6 +37,7 @@ except Exception:
 
 
 if Required is None and _AnnotatedAlias is None:
+    # No extras available, so no `include_extras`
     get_type_hints = _get_type_hints
 else:
 
@@ -29,6 +46,7 @@ else:
 
 
 def get_typeddict_hints(obj):
+    """Same as `get_type_hints`, but strips off Required/NotRequired"""
     hints = get_type_hints(obj)
     out = {}
     for k, v in hints.items():
@@ -37,3 +55,478 @@ def get_typeddict_hints(obj):
             v = v.__args__[0]
         out[k] = v
     return out
+
+
+def schema(type: Any) -> Dict[str, Any]:
+    """Generate a JSON Schema for a given type.
+
+    Any schemas for (potentially) shared components are extracted and stored in
+    a top-level ``"$defs"`` field.
+
+    If you want to generate schemas for multiple types, or to have more control
+    over the generated schema you may want to use ``schema_components`` instead.
+
+    Parameters
+    ----------
+    type : Type
+        The type to generate the schema for.
+
+    Returns
+    -------
+    schema : dict
+        The generated JSON Schema.
+
+    See Also
+    --------
+    schema_components
+    """
+    (out,), components = schema_components((type,))
+    if components:
+        out["$defs"] = components
+    return out
+
+
+def schema_components(
+    types: Iterable[Any], ref_template: str = "#/$defs/{name}"
+) -> Tuple[Tuple[Dict[str, Any], ...], Dict[str, Any]]:
+    """Generate JSON Schemas for one or more types.
+
+    Any schemas for (potentially) shared components are extracted and returned
+    in a separate ``components`` dict.
+
+    Parameters
+    ----------
+    types : Iterable[Type]
+        An iterable of one or more types to generate schemas for.
+    ref_template : str, optional
+        A template to use when generating ``"$ref"`` fields. This template is
+        formatted with the type name as ``template.format(name=name)``. This
+        can be useful if you intend to store the ``components`` mapping
+        somewhere other than a top-level ``"$defs"`` field. For example, you
+        might use ``ref_template="#/components/{name}"`` if generating an
+        OpenAPI schema.
+
+    Returns
+    -------
+    schemas : tuple[dict]
+        A tuple of JSON Schemas, one for each type in ``types``.
+    components : dict
+        A mapping of name to schema for any shared components used by
+        ``schemas``.
+
+    See Also
+    --------
+    schema
+    """
+    return SchemaBuilder(types, ref_template).run()
+
+
+def roundtrip_json(d):
+    from ._core import json_encode, json_decode
+
+    return json_decode(json_encode(d))
+
+
+def merge_json(a, b):
+    if b:
+        a = a.copy()
+        for key, b_val in b.items():
+            if key in a:
+                a_val = a[key]
+                if isinstance(a_val, dict) and isinstance(b_val, dict):
+                    a[key] = merge_json(a_val, b_val)
+                elif isinstance(a_val, (list, tuple)) and isinstance(
+                    b_val, (list, tuple)
+                ):
+                    a[key] = list(a_val) + list(b_val)
+                else:
+                    a[key] = b_val
+            else:
+                a[key] = b_val
+    return a
+
+
+def is_struct(t):
+    from ._core import Struct
+
+    return type(t) is type(Struct)
+
+
+def is_enum(t):
+    return type(t) is enum.EnumMeta
+
+
+def is_typeddict(t):
+    try:
+        return issubclass(t, dict) and hasattr(t, "__total__")
+    except TypeError:
+        return False
+
+
+def is_namedtuple(t):
+    try:
+        return issubclass(t, tuple) and hasattr(t, "_fields")
+    except TypeError:
+        return False
+
+
+def has_nondefault_docstring(t):
+    if not (doc := getattr(t, "__doc__", None)):
+        return False
+
+    if is_enum(t):
+        # TODO: The startswith check was added for compatibility with the
+        # Python 3.11 release candidate, but the cpython code requiring this
+        # change has since been removed. This check can be removed once Python
+        # 3.11 is released.
+        return not (
+            doc == "An enumeration."
+            or doc.startswith("A collection of name/value pairs.")
+        )
+    elif is_namedtuple(t):
+        return not (doc.startswith(f"{t.__name__}(") and doc.endswith(")"))
+    return True
+
+
+UNSET = type("Unset", (), {"__repr__": lambda s: "UNSET"})()
+
+
+def origin_args_metadata(t):
+    from ._core import Meta
+
+    if type(t) is _AnnotatedAlias:
+        metadata = tuple(m for m in t.__metadata__ if type(m) is Meta)
+        t = t.__origin__
+    else:
+        metadata = ()
+
+    if type(t) is UnionType:
+        args = t.__args__
+        t = Union
+    elif t in (List, Tuple, Set, FrozenSet, Dict):
+        t = t.__origin__
+        args = None
+    elif hasattr(t, "__origin__"):
+        args = getattr(t, "__args__", None)
+        t = t.__origin__
+    else:
+        args = None
+
+    return t, args, metadata
+
+
+class SchemaBuilder:
+    def __init__(self, types, ref_template):
+        self.types = tuple(types)
+        self.ref_template = ref_template
+        # Cache of type hints per type
+        self.type_hints = {}
+        # Collections of component types to extract
+        self.structs = set()
+        self.enums = set()
+        self.typeddicts = set()
+        self.namedtuples = set()
+
+    def _get_type_hints(self, t: Any) -> dict:
+        """A cached version of `get_type_hints`"""
+        try:
+            return self.type_hints[t]
+        except KeyError:
+            out = self.type_hints[t] = get_type_hints(t)
+            return out
+
+    @property
+    def subtypes(self):
+        return (*self.structs, *self.enums, *self.typeddicts, *self.namedtuples)
+
+    def run(self):
+        # First construct a decoder to validate the types are valid
+        from ._core import JSONDecoder
+
+        JSONDecoder(Tuple[self.types])
+
+        for t in self.types:
+            self._collect_type(t)
+
+        self._init_name_map()
+
+        schemas = [self._type_to_schema(t) for t in self.types]
+        components = {
+            self.subtype_names[t]: self._type_to_schema(t, check_ref=False)
+            for t in self.subtypes
+        }
+        return schemas, components
+
+    def _collect_type(self, t):
+        t, args, metadata = origin_args_metadata(t)
+
+        if is_enum(t):
+            self.enums.add(t)
+        elif is_struct(t):
+            if t not in self.structs:
+                self.structs.add(t)
+                self._collect_fields(t)
+        elif is_typeddict(t):
+            if t not in self.typeddicts:
+                self.typeddicts.add(t)
+                self._collect_fields(t)
+        elif is_namedtuple(t):
+            if t not in self.namedtuples:
+                self.namedtuples.add(t)
+                self._collect_fields(t)
+        elif t in (dict, list, tuple, set, frozenset, Union) and args:
+            for a in args:
+                if a is not ...:
+                    self._collect_type(a)
+
+    def _collect_fields(self, t):
+        fields = self._get_type_hints(t)
+        for v in fields.values():
+            self._collect_type(v)
+
+    def _init_name_map(self):
+        def normalize(name):
+            return re.sub(r"[^a-zA-Z0-9.\-_]", "_", name)
+
+        def fullname(t):
+            return normalize(f"{t.__module__}.{t.__qualname__}")
+
+        conflicts = set()
+        names = {}
+
+        for t in self.subtypes:
+            name = normalize(t.__name__)
+            if name in names:
+                old_t = names.pop(name)
+                conflicts.add(name)
+                names[fullname(old_t)] = old_t
+            if name in conflicts:
+                names[fullname(t)] = t
+            else:
+                names[name] = t
+        self.subtype_names = {v: k for k, v in names.items()}
+
+    def _process_metadata(self, t, metadata):
+        schema = {}
+        extra = {}
+        for meta in metadata:
+            if meta.title:
+                schema["title"] = meta.title
+            if meta.description:
+                schema["description"] = meta.description
+            if meta.examples:
+                schema["examples"] = roundtrip_json(meta.examples)
+            if meta.extra_json_schema is not None:
+                extra = merge_json(extra, roundtrip_json(meta.extra_json_schema))
+            if meta.ge is not None:
+                schema["minimum"] = meta.ge
+            if meta.gt is not None:
+                schema["exclusiveMinimum"] = meta.gt
+            if meta.le is not None:
+                schema["maximum"] = meta.le
+            if meta.lt is not None:
+                schema["exclusiveMaximum"] = meta.lt
+            if meta.multiple_of is not None:
+                schema["multipleOf"] = meta.multiple_of
+            if meta.pattern is not None:
+                schema["pattern"] = meta.pattern
+            if t is str:
+                if meta.max_length is not None:
+                    schema["maxLength"] = meta.max_length
+                if meta.min_length is not None:
+                    schema["minLength"] = meta.min_length
+            elif t in (bytes, bytearray):
+                if meta.max_length is not None:
+                    schema["maxLength"] = 4 * ((meta.max_length + 2) // 3)
+                if meta.min_length is not None:
+                    schema["minLength"] = 4 * ((meta.min_length + 2) // 3)
+            elif t in (list, tuple, set, frozenset):
+                if meta.max_length is not None:
+                    schema["maxItems"] = meta.max_length
+                if meta.min_length is not None:
+                    schema["minItems"] = meta.min_length
+            elif t is dict:
+                if meta.max_length is not None:
+                    schema["maxProperties"] = meta.max_length
+                if meta.min_length is not None:
+                    schema["minProperties"] = meta.min_length
+        return schema, extra
+
+    def _type_to_schema(self, typ, *, default=UNSET, check_ref=True):
+        from ._core import Raw, Ext
+
+        t, args, metadata = origin_args_metadata(typ)
+        schema, extra = self._process_metadata(t, metadata)
+
+        if check_ref and (name := self.subtype_names.get(t)):
+            schema["$ref"] = self.ref_template.format(name=name)
+            return merge_json(schema, extra)
+
+        if default is not UNSET:
+            schema["default"] = roundtrip_json(default)
+
+        if t in (Any, Raw):
+            pass
+        elif t is None or t is type(None):
+            schema["type"] = "null"
+        elif t is bool:
+            schema["type"] = "boolean"
+        elif t is int:
+            schema["type"] = "integer"
+        elif t is float:
+            schema["type"] = "number"
+        elif t is str:
+            schema["type"] = "string"
+        elif t in (bytes, bytearray):
+            schema["type"] = "string"
+            schema["contentEncoding"] = "base64"
+        elif t is datetime.datetime:
+            schema["type"] = "string"
+            schema["format"] = "date-time"
+        elif t in (list, set, frozenset):
+            schema["type"] = "array"
+            if args:
+                schema["items"] = self._type_to_schema(args[0])
+        elif t is tuple:
+            schema["type"] = "array"
+            if args is not None:
+                # Handle an annoying compatibility issue:
+                # - Tuple[()] has args == ((),)
+                # - tuple[()] has args == ()
+                if args == ((),):
+                    args = ()
+                if len(args) == 2 and args[-1] is ...:
+                    schema["items"] = self._type_to_schema(args[0])
+                elif args:
+                    schema["prefixItems"] = [self._type_to_schema(a) for a in args]
+                    schema["minItems"] = len(args)
+                    schema["maxItems"] = len(args)
+                    schema["items"] = False
+                else:
+                    schema["minItems"] = 0
+                    schema["maxItems"] = 0
+        elif t is dict:
+            schema["type"] = "object"
+            if args:
+                schema["additionalProperties"] = self._type_to_schema(args[1])
+        elif t is Union:
+            structs = {}
+            other = []
+            tag_field = None
+            for a in args:
+                a_t, a_args, a_metadata = origin_args_metadata(a)
+                if is_struct(a_t) and not a_t.array_like:
+                    tag_field = a_t.__struct_tag_field__
+                    structs[a_t.__struct_tag__] = a
+                else:
+                    other.append(a)
+
+            options = [self._type_to_schema(a) for a in other]
+
+            if len(structs) >= 2:
+                mapping = {
+                    k: self.ref_template.format(name=self.subtype_names[v])
+                    for k, v in structs.items()
+                }
+                struct_schema = {
+                    "anyOf": [self._type_to_schema(s) for s in structs.values()],
+                    "discriminator": {"propertyName": tag_field, "mapping": mapping},
+                }
+                if options:
+                    options.append(struct_schema)
+                    schema["anyOf"] = options
+                else:
+                    schema.update(struct_schema)
+            else:
+                if len(structs) == 1:
+                    options.append(self._type_to_schema(*structs.values()))
+                schema["anyOf"] = options
+        elif t is Literal:
+            schema["enum"] = sorted(args)
+        elif is_enum(t):
+            schema.setdefault("title", t.__name__)
+            if has_nondefault_docstring(t):
+                schema.setdefault("description", t.__doc__)
+            if issubclass(t, enum.IntEnum):
+                schema["enum"] = sorted(int(e) for e in t)
+            else:
+                schema["enum"] = sorted(e.name for e in t)
+        elif is_struct(t):
+            schema.setdefault("title", t.__name__)
+            if has_nondefault_docstring(t):
+                schema.setdefault("description", t.__doc__)
+            hints = self._get_type_hints(t)
+            required = []
+            names = []
+            fields = []
+
+            if t.__struct_tag_field__:
+                required.append(t.__struct_tag_field__)
+                names.append(t.__struct_tag_field__)
+                fields.append({"enum": [t.__struct_tag__]})
+
+            n_required = len(t.__struct_fields__) - len(t.__struct_defaults__)
+            required.extend(t.__struct_encode_fields__[:n_required])
+            names.extend(t.__struct_encode_fields__)
+            fields.extend(
+                self._type_to_schema(hints[f], default=d)
+                for f, ef, d in zip(
+                    t.__struct_fields__,
+                    t.__struct_encode_fields__,
+                    (UNSET,) * n_required + t.__struct_defaults__,
+                )
+            )
+            if t.array_like:
+                schema["type"] = "array"
+                schema["prefixItems"] = fields
+                schema["minItems"] = len(required)
+            else:
+                schema["type"] = "object"
+                schema["properties"] = dict(zip(names, fields))
+                schema["required"] = required
+        elif is_typeddict(t):
+            schema.setdefault("title", t.__name__)
+            if has_nondefault_docstring(t):
+                schema.setdefault("description", t.__doc__)
+            hints = self._get_type_hints(t)
+            schema["type"] = "object"
+            schema["properties"] = {
+                field: self._type_to_schema(field_type)
+                for field, field_type in hints.items()
+            }
+            if hasattr(t, "__required_keys__"):
+                required = sorted(t.__required_keys__)
+            elif t.__total__:
+                required = sorted(hints)
+            else:
+                required = []
+            schema["required"] = required
+        elif is_namedtuple(t):
+            schema.setdefault("title", t.__name__)
+            if has_nondefault_docstring(t):
+                schema.setdefault("description", t.__doc__)
+            hints = self._get_type_hints(t)
+            fields = [
+                self._type_to_schema(
+                    hints.get(f, Any), default=t._field_defaults.get(f, UNSET)
+                )
+                for f in t._fields
+            ]
+            schema["type"] = "array"
+            schema["prefixItems"] = fields
+            schema["minItems"] = len(t._fields) - len(t._field_defaults)
+            schema["maxItems"] = len(t._fields)
+        elif t is Ext:
+            raise TypeError("json-schema doesn't support msgpack Ext types")
+        elif not extra:
+            # `t` is a custom type, an explicit schema is required
+            raise TypeError(
+                "Custom types currently require a schema be explicitly provided "
+                "by annotating the type with `Meta(extra_json_schema=...)` - type "
+                f"{t!r} is not supported"
+            )
+
+        if extra:
+            schema = merge_json(schema, extra)
+
+        return schema
diff --git a/msgspec/json.py b/msgspec/json.py
index 4808e68..410ac00 100644
--- a/msgspec/json.py
+++ b/msgspec/json.py
@@ -4,3 +4,4 @@ from ._core import (
     json_encode as encode,
     json_decode as decode,
 )
+from ._utils import schema, schema_components
diff --git a/msgspec/json.pyi b/msgspec/json.pyi
index 7f26768..48b49e7 100644
--- a/msgspec/json.pyi
+++ b/msgspec/json.pyi
@@ -1,4 +1,15 @@
-from typing import Any, Type, TypeVar, Generic, Optional, Callable, overload
+from typing import (
+    Any,
+    Type,
+    Tuple,
+    Dict,
+    TypeVar,
+    Generic,
+    Optional,
+    Callable,
+    overload,
+)
+from collections.abc import Iterable
 
 T = TypeVar("T")
 
@@ -51,3 +62,7 @@ def decode(
     dec_hook: dec_hook_sig = None,
 ) -> T: ...
 def encode(obj: Any, *, enc_hook: enc_hook_sig = None) -> bytes: ...
+def schema(type: Any) -> Dict[str, Any]: ...
+def schema_components(
+    types: Iterable[Any], ref_template: str = "#/$defs/{name}"
+) -> Tuple[Tuple[Dict[str, Any], ...], Dict[str, Any]]: ...


Test patch:
diff --git a/tests/basic_typing_examples.py b/tests/basic_typing_examples.py
index ba56302..c939e92 100644
--- a/tests/basic_typing_examples.py
+++ b/tests/basic_typing_examples.py
@@ -505,3 +505,25 @@ def check_json_decode_dec_hook() -> None:
 
     msgspec.json.decode(b"test", dec_hook=dec_hook)
     msgspec.json.Decoder(dec_hook=dec_hook)
+
+
+##########################################################
+# JSON Schema                                            #
+##########################################################
+
+
+def check_json_schema() -> None:
+    o = msgspec.json.schema(List[int])
+    reveal_type(o)  # assert ("Dict" in typ or "dict" in typ)
+
+
+def check_json_schema_components() -> None:
+    s, c = msgspec.json.schema_components([List[int]])
+    reveal_type(s)  # assert ("dict" in typ.lower()) and ("tuple" in typ.lower())
+    reveal_type(c)  # assert ("dict" in typ.lower())
+
+
+def check_json_schema_components_full() -> None:
+    s, c = msgspec.json.schema_components([List[int]], ref_template="#/definitions/{name}")
+    reveal_type(s)  # assert ("dict" in typ.lower()) and ("tuple" in typ.lower())
+    reveal_type(c)  # assert ("dict" in typ.lower())
diff --git a/tests/test_schema.py b/tests/test_schema.py
new file mode 100644
index 0000000..86fc520
--- /dev/null
+++ b/tests/test_schema.py
@@ -0,0 +1,827 @@
+import enum
+import datetime
+import uuid
+from base64 import b64encode
+from copy import deepcopy
+from collections import namedtuple
+from typing import (
+    Any,
+    Dict,
+    FrozenSet,
+    List,
+    Literal,
+    NamedTuple,
+    Set,
+    Tuple,
+    TypedDict,
+    Union,
+)
+
+import pytest
+
+import msgspec
+from msgspec._utils import merge_json
+
+from utils import temp_module
+
+try:
+    from typing import Annotated
+except ImportError:
+    try:
+        from typing_extensions import Annotated
+    except ImportError:
+        pytestmark = pytest.mark.skip("Annotated types not available")
+
+
+@pytest.mark.parametrize(
+    "a,b,sol",
+    [
+        (
+            {"a": {"b": {"c": 1}}},
+            {"a": {"b": {"d": 2}}},
+            {"a": {"b": {"c": 1, "d": 2}}},
+        ),
+        ({"a": {"b": {"c": 1}}}, {"a": {"b": 2}}, {"a": {"b": 2}}),
+        ({"a": [1, 2]}, {"a": [3, 4]}, {"a": [1, 2, 3, 4]}),
+        ({"a": {"b": 1}}, {"a2": 3}, {"a": {"b": 1}, "a2": 3}),
+        ({"a": 1}, {}, {"a": 1}),
+    ],
+)
+def test_merge_json(a, b, sol):
+    a_orig = deepcopy(a)
+    b_orig = deepcopy(b)
+    res = merge_json(a, b)
+    assert res == sol
+    assert a == a_orig
+    assert b == b_orig
+
+
+def type_index(typ, args):
+    try:
+        return typ[args]
+    except TypeError:
+        pytest.skip("Not supported in Python 3.8")
+
+
+def test_any():
+    assert msgspec.json.schema(Any) == {}
+
+
+def test_raw():
+    assert msgspec.json.schema(msgspec.Raw) == {}
+
+
+def test_msgpack_ext():
+    with pytest.raises(TypeError):
+        assert msgspec.json.schema(msgspec.msgpack.Ext)
+
+
+def test_custom():
+    with pytest.raises(TypeError, match="Custom types"):
+        assert msgspec.json.schema(uuid.UUID)
+
+    schema = {"type": "string", "pattern": "uuid4"}
+
+    assert (
+        msgspec.json.schema(
+            Annotated[uuid.UUID, msgspec.Meta(extra_json_schema=schema)]
+        )
+        == schema
+    )
+
+
+def test_none():
+    assert msgspec.json.schema(None) == {"type": "null"}
+
+
+def test_bool():
+    assert msgspec.json.schema(bool) == {"type": "boolean"}
+
+
+def test_int():
+    assert msgspec.json.schema(int) == {"type": "integer"}
+
+
+def test_float():
+    assert msgspec.json.schema(float) == {"type": "number"}
+
+
+def test_string():
+    assert msgspec.json.schema(str) == {"type": "string"}
+
+
+@pytest.mark.parametrize("typ", [bytes, bytearray])
+def test_binary(typ):
+    assert msgspec.json.schema(typ) == {
+        "type": "string",
+        "contentEncoding": "base64",
+    }
+
+
+def test_datetime():
+    assert msgspec.json.schema(datetime.datetime) == {
+        "type": "string",
+        "format": "date-time",
+    }
+
+
+@pytest.mark.parametrize(
+    "typ", [list, tuple, set, frozenset, List, Tuple, Set, FrozenSet]
+)
+def test_sequence_any(typ):
+    assert msgspec.json.schema(typ) == {"type": "array"}
+
+
+@pytest.mark.parametrize(
+    "cls", [list, tuple, set, frozenset, List, Tuple, Set, FrozenSet]
+)
+def test_sequence_typed(cls):
+    args = (int, ...) if cls in (tuple, Tuple) else int
+    typ = type_index(cls, args)
+    assert msgspec.json.schema(typ) == {"type": "array", "items": {"type": "integer"}}
+
+
+@pytest.mark.parametrize("cls", [tuple, Tuple])
+def test_tuple(cls):
+    typ = type_index(cls, (int, float, str))
+    assert msgspec.json.schema(typ) == {
+        "type": "array",
+        "minItems": 3,
+        "maxItems": 3,
+        "items": False,
+        "prefixItems": [
+            {"type": "integer"},
+            {"type": "number"},
+            {"type": "string"},
+        ],
+    }
+
+
+@pytest.mark.parametrize("cls", [tuple, Tuple])
+def test_empty_tuple(cls):
+    typ = type_index(cls, ())
+    assert msgspec.json.schema(typ) == {
+        "type": "array",
+        "minItems": 0,
+        "maxItems": 0,
+    }
+
+
+@pytest.mark.parametrize("typ", [dict, Dict])
+def test_dict_any(typ):
+    assert msgspec.json.schema(typ) == {"type": "object"}
+
+
+@pytest.mark.parametrize("cls", [dict, Dict])
+def test_dict_typed(cls):
+    typ = type_index(cls, (str, int))
+    assert msgspec.json.schema(typ) == {
+        "type": "object",
+        "additionalProperties": {"type": "integer"},
+    }
+
+
+def test_int_enum():
+    class Example(enum.IntEnum):
+        C = 1
+        B = 3
+        A = 2
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {"Example": {"title": "Example", "enum": [1, 2, 3]}},
+    }
+
+
+def test_enum():
+    class Example(enum.Enum):
+        """A docstring"""
+
+        C = "x"
+        B = "z"
+        A = "y"
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "description": "A docstring",
+                "enum": ["A", "B", "C"],
+            }
+        },
+    }
+
+
+def test_int_literal():
+    assert msgspec.json.schema(Literal[3, 1, 2]) == {"enum": [1, 2, 3]}
+
+
+def test_str_literal():
+    assert msgspec.json.schema(Literal["c", "a", "b"]) == {"enum": ["a", "b", "c"]}
+
+
+def test_struct_object():
+    class Point(msgspec.Struct):
+        x: int
+        y: int
+
+    class Polygon(msgspec.Struct):
+        """An example docstring"""
+
+        vertices: List[Point]
+        name: Union[str, None] = None
+        metadata: Dict[str, str] = {}
+
+    assert msgspec.json.schema(Polygon) == {
+        "$ref": "#/$defs/Polygon",
+        "$defs": {
+            "Polygon": {
+                "title": "Polygon",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "vertices": {
+                        "type": "array",
+                        "items": {"$ref": "#/$defs/Point"},
+                    },
+                    "name": {
+                        "anyOf": [{"type": "string"}, {"type": "null"}],
+                        "default": None,
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": {"type": "string"},
+                        "default": {},
+                    },
+                },
+                "required": ["vertices"],
+            },
+            "Point": {
+                "title": "Point",
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["x", "y"],
+            },
+        },
+    }
+
+
+def test_struct_array_like():
+    class Example(msgspec.Struct, array_like=True):
+        """An example docstring"""
+
+        a: int
+        b: str
+        c: List[int] = []
+        d: Dict[str, int] = {}
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "description": "An example docstring",
+                "type": "array",
+                "prefixItems": [
+                    {"type": "integer"},
+                    {"type": "string"},
+                    {"type": "array", "items": {"type": "integer"}, "default": []},
+                    {
+                        "type": "object",
+                        "additionalProperties": {"type": "integer"},
+                        "default": {},
+                    },
+                ],
+                "minItems": 2,
+            }
+        },
+    }
+
+
+def test_struct_no_fields():
+    class Example(msgspec.Struct):
+        pass
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "type": "object",
+                "properties": {},
+                "required": [],
+            }
+        },
+    }
+
+
+def test_struct_object_tagged():
+    class Point(msgspec.Struct, tag=True):
+        x: int
+        y: int
+
+    assert msgspec.json.schema(Point) == {
+        "$ref": "#/$defs/Point",
+        "$defs": {
+            "Point": {
+                "title": "Point",
+                "type": "object",
+                "properties": {
+                    "type": {"enum": ["Point"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["type", "x", "y"],
+            }
+        },
+    }
+
+
+def test_struct_array_tagged():
+    class Point(msgspec.Struct, tag=True, array_like=True):
+        x: int
+        y: int
+
+    assert msgspec.json.schema(Point) == {
+        "$ref": "#/$defs/Point",
+        "$defs": {
+            "Point": {
+                "title": "Point",
+                "type": "array",
+                "prefixItems": [
+                    {"enum": ["Point"]},
+                    {"type": "integer"},
+                    {"type": "integer"},
+                ],
+                "minItems": 3,
+            }
+        },
+    }
+
+
+def test_typing_namedtuple():
+    class Example(NamedTuple):
+        """An example docstring"""
+
+        a: str
+        b: bool
+        c: int = 0
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "description": "An example docstring",
+                "type": "array",
+                "prefixItems": [
+                    {"type": "string"},
+                    {"type": "boolean"},
+                    {"type": "integer", "default": 0},
+                ],
+                "minItems": 2,
+                "maxItems": 3,
+            }
+        },
+    }
+
+
+def test_collections_namedtuple():
+    Example = namedtuple("Example", ["a", "b", "c"], defaults=(0,))
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "type": "array",
+                "prefixItems": [{}, {}, {"default": 0}],
+                "minItems": 2,
+                "maxItems": 3,
+            }
+        },
+    }
+
+
+@pytest.mark.parametrize("use_typing_extensions", [False, True])
+def test_typeddict(use_typing_extensions):
+    if use_typing_extensions:
+        tex = pytest.importorskip("typing_extensions")
+        cls = tex.TypedDict
+    else:
+        cls = TypedDict
+
+    class Example(cls):
+        """An example docstring"""
+
+        a: str
+        b: bool
+        c: int
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "a": {"type": "string"},
+                    "b": {"type": "boolean"},
+                    "c": {"type": "integer"},
+                },
+                "required": ["a", "b", "c"],
+            }
+        },
+    }
+
+
+@pytest.mark.parametrize("use_typing_extensions", [False, True])
+def test_typeddict_optional(use_typing_extensions):
+    if use_typing_extensions:
+        tex = pytest.importorskip("typing_extensions")
+        cls = tex.TypedDict
+    else:
+        cls = TypedDict
+
+    class Base(cls):
+        a: str
+        b: bool
+
+    class Example(Base, total=False):
+        """An example docstring"""
+
+        c: int
+
+    if not hasattr(Example, "__required_keys__"):
+        # This should be Python 3.8, builtin typing only
+        pytest.skip("partially optional TypedDict not supported")
+
+    assert msgspec.json.schema(Example) == {
+        "$ref": "#/$defs/Example",
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "a": {"type": "string"},
+                    "b": {"type": "boolean"},
+                    "c": {"type": "integer"},
+                },
+                "required": ["a", "b"],
+            }
+        },
+    }
+
+
+@pytest.mark.parametrize("use_union_operator", [False, True])
+def test_union(use_union_operator):
+    class Example(msgspec.Struct):
+        x: int
+        y: int
+
+    if use_union_operator:
+        try:
+            typ = int | str | Example
+        except TypeError:
+            pytest.skip("Union operator not supported")
+    else:
+        typ = Union[int, str, Example]
+
+    assert msgspec.json.schema(typ) == {
+        "anyOf": [
+            {"type": "integer"},
+            {"type": "string"},
+            {"$ref": "#/$defs/Example"},
+        ],
+        "$defs": {
+            "Example": {
+                "title": "Example",
+                "type": "object",
+                "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}},
+                "required": ["x", "y"],
+            }
+        },
+    }
+
+
+def test_struct_tagged_union():
+    class Point(msgspec.Struct, tag=True):
+        x: int
+        y: int
+
+    class Point3D(Point):
+        z: int
+
+    assert msgspec.json.schema(Union[Point, Point3D]) == {
+        "anyOf": [{"$ref": "#/$defs/Point"}, {"$ref": "#/$defs/Point3D"}],
+        "discriminator": {
+            "mapping": {"Point": "#/$defs/Point", "Point3D": "#/$defs/Point3D"},
+            "propertyName": "type",
+        },
+        "$defs": {
+            "Point": {
+                "properties": {
+                    "type": {"enum": ["Point"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["type", "x", "y"],
+                "title": "Point",
+                "type": "object",
+            },
+            "Point3D": {
+                "properties": {
+                    "type": {"enum": ["Point3D"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                    "z": {"type": "integer"},
+                },
+                "required": ["type", "x", "y", "z"],
+                "title": "Point3D",
+                "type": "object",
+            },
+        },
+    }
+
+
+def test_struct_tagged_union_mixed_types():
+    class Point(msgspec.Struct, tag=True):
+        x: int
+        y: int
+
+    class Point3D(Point):
+        z: int
+
+    assert msgspec.json.schema(Union[Point, Point3D, int, float]) == {
+        "anyOf": [
+            {"type": "integer"},
+            {"type": "number"},
+            {
+                "anyOf": [{"$ref": "#/$defs/Point"}, {"$ref": "#/$defs/Point3D"}],
+                "discriminator": {
+                    "mapping": {"Point": "#/$defs/Point", "Point3D": "#/$defs/Point3D"},
+                    "propertyName": "type",
+                },
+            },
+        ],
+        "$defs": {
+            "Point": {
+                "properties": {
+                    "type": {"enum": ["Point"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["type", "x", "y"],
+                "title": "Point",
+                "type": "object",
+            },
+            "Point3D": {
+                "properties": {
+                    "type": {"enum": ["Point3D"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                    "z": {"type": "integer"},
+                },
+                "required": ["type", "x", "y", "z"],
+                "title": "Point3D",
+                "type": "object",
+            },
+        },
+    }
+
+
+def test_struct_array_union():
+    class Point(msgspec.Struct, array_like=True, tag=True):
+        x: int
+        y: int
+
+    class Point3D(Point):
+        z: int
+
+    assert msgspec.json.schema(Union[Point, Point3D]) == {
+        "anyOf": [{"$ref": "#/$defs/Point"}, {"$ref": "#/$defs/Point3D"}],
+        "$defs": {
+            "Point": {
+                "minItems": 3,
+                "prefixItems": [
+                    {"enum": ["Point"]},
+                    {"type": "integer"},
+                    {"type": "integer"},
+                ],
+                "title": "Point",
+                "type": "array",
+            },
+            "Point3D": {
+                "minItems": 4,
+                "prefixItems": [
+                    {"enum": ["Point3D"]},
+                    {"type": "integer"},
+                    {"type": "integer"},
+                    {"type": "integer"},
+                ],
+                "title": "Point3D",
+                "type": "array",
+            },
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    "field, constraint",
+    [
+        ("ge", "minimum"),
+        ("gt", "exclusiveMinimum"),
+        ("le", "maximum"),
+        ("lt", "exclusiveMaximum"),
+        ("multiple_of", "multipleOf"),
+    ],
+)
+def test_numeric_metadata(field, constraint):
+    typ = Annotated[int, msgspec.Meta(**{field: 2})]
+    assert msgspec.json.schema(typ) == {"type": "integer", constraint: 2}
+
+
+@pytest.mark.parametrize(
+    "field, val, constraint",
+    [
+        ("pattern", "[a-z]*", "pattern"),
+        ("min_length", 0, "minLength"),
+        ("max_length", 3, "maxLength"),
+    ],
+)
+def test_string_metadata(field, val, constraint):
+    typ = Annotated[str, msgspec.Meta(**{field: val})]
+    assert msgspec.json.schema(typ) == {"type": "string", constraint: val}
+
+
+@pytest.mark.parametrize("typ", [bytes, bytearray])
+@pytest.mark.parametrize(
+    "field, n, constraint",
+    [("min_length", 2, "minLength"), ("max_length", 7, "maxLength")],
+)
+def test_binary_metadata(typ, field, n, constraint):
+    n2 = len(b64encode(b"x" * n))
+    typ = Annotated[typ, msgspec.Meta(**{field: n})]
+    assert msgspec.json.schema(typ) == {
+        "type": "string",
+        constraint: n2,
+        "contentEncoding": "base64",
+    }
+
+
+@pytest.mark.parametrize("typ", [list, tuple, set, frozenset])
+@pytest.mark.parametrize(
+    "field, constraint",
+    [("min_length", "minItems"), ("max_length", "maxItems")],
+)
+def test_array_metadata(typ, field, constraint):
+    typ = Annotated[typ, msgspec.Meta(**{field: 2})]
+    assert msgspec.json.schema(typ) == {"type": "array", constraint: 2}
+
+
+@pytest.mark.parametrize(
+    "field, constraint",
+    [("min_length", "minProperties"), ("max_length", "maxProperties")],
+)
+def test_object_metadata(field, constraint):
+    typ = Annotated[dict, msgspec.Meta(**{field: 2})]
+    assert msgspec.json.schema(typ) == {"type": "object", constraint: 2}
+
+
+def test_generic_metadata():
+    typ = Annotated[
+        int,
+        msgspec.Meta(
+            title="the title",
+            description="the description",
+            examples=[1, 2, 3],
+            extra_json_schema={"title": "an override", "default": 1},
+        ),
+    ]
+    assert msgspec.json.schema(typ) == {
+        "type": "integer",
+        "title": "an override",
+        "description": "the description",
+        "examples": [1, 2, 3],
+        "default": 1,
+    }
+
+
+def test_component_names_collide():
+    s1 = """
+    import msgspec
+    Ex = msgspec.defstruct("Ex", [("x", int), ("y", int)])
+    """
+
+    s2 = """
+    import msgspec
+    Ex = msgspec.defstruct("Ex", [("a", str), ("b", str)])
+    """
+
+    with temp_module(s1) as m1, temp_module(s2) as m2:
+        (r1, r2), components = msgspec.json.schema_components([m1.Ex, m2.Ex])
+
+    assert r1 == {"$ref": f"#/$defs/{m1.__name__}.Ex"}
+    assert r2 == {"$ref": f"#/$defs/{m2.__name__}.Ex"}
+    assert components == {
+        f"{m1.__name__}.Ex": {
+            "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}},
+            "required": ["x", "y"],
+            "title": "Ex",
+            "type": "object",
+        },
+        f"{m2.__name__}.Ex": {
+            "properties": {"a": {"type": "string"}, "b": {"type": "string"}},
+            "required": ["a", "b"],
+            "title": "Ex",
+            "type": "object",
+        },
+    }
+
+
+def test_schema_components_collects_subtypes():
+    class ExEnum(enum.Enum):
+        A = 1
+
+    class ExStruct(msgspec.Struct):
+        b: Union[Set[FrozenSet[ExEnum]], int]
+
+    class ExDict(TypedDict):
+        c: Tuple[ExStruct, ...]
+
+    class ExTuple(NamedTuple):
+        d: List[ExDict]
+
+    (s,), components = msgspec.json.schema_components([Dict[str, ExTuple]])
+
+    r1 = {"$ref": "#/$defs/ExEnum"}
+    r2 = {"$ref": "#/$defs/ExStruct"}
+    r3 = {"$ref": "#/$defs/ExDict"}
+    r4 = {"$ref": "#/$defs/ExTuple"}
+
+    assert s == {"type": "object", "additionalProperties": r4}
+    assert components == {
+        "ExEnum": {"enum": ["A"], "title": "ExEnum"},
+        "ExStruct": {
+            "type": "object",
+            "title": "ExStruct",
+            "properties": {
+                "b": {
+                    "anyOf": [
+                        {"items": {"items": r1, "type": "array"}, "type": "array"},
+                        {"type": "integer"},
+                    ]
+                }
+            },
+            "required": ["b"],
+        },
+        "ExDict": {
+            "title": "ExDict",
+            "type": "object",
+            "properties": {"c": {"items": r2, "type": "array"}},
+            "required": ["c"],
+        },
+        "ExTuple": {
+            "title": "ExTuple",
+            "type": "array",
+            "prefixItems": [{"items": r3, "type": "array"}],
+            "maxItems": 1,
+            "minItems": 1,
+        },
+    }
+
+
+def test_ref_template():
+    class Ex1(msgspec.Struct):
+        a: int
+
+    class Ex2(msgspec.Struct):
+        b: Ex1
+
+    (s1, s2), components = msgspec.json.schema_components(
+        [Ex1, Ex2], ref_template="#/definitions/{name}"
+    )
+
+    assert s1 == {"$ref": "#/definitions/Ex1"}
+    assert s2 == {"$ref": "#/definitions/Ex2"}
+
+    assert components == {
+        "Ex1": {
+            "title": "Ex1",
+            "type": "object",
+            "properties": {"a": {"type": "integer"}},
+            "required": ["a"],
+        },
+        "Ex2": {
+            "title": "Ex2",
+            "type": "object",
+            "properties": {"b": s1},
+            "required": ["b"],
+        },
+    }