Repository: jcrist/msgspec
Pull request #393 (2023-04-25 05:22:11)

Support Generic structs in `msgspec.inspect`/`msgspec.json.schema`
This adds support for generic structs to our introspection utilities (`msgspec.inspect`/`msgspec.json.schema`). This is a follow-up to #386.


Patch:
diff --git a/msgspec/_json_schema.py b/msgspec/_json_schema.py
index 5f6cbb0..215dd0c 100644
--- a/msgspec/_json_schema.py
+++ b/msgspec/_json_schema.py
@@ -124,6 +124,35 @@ def _collect_component_types(type_infos: Iterable[mi.Type]) -> dict[Any, mi.Type
     return components
 
 
+def _type_repr(obj):
+    return obj.__name__ if isinstance(obj, type) else repr(obj)
+
+
+def _get_class_name(cls: Any) -> str:
+    if hasattr(cls, "__origin__"):
+        name = cls.__origin__.__name__
+        args = ", ".join(_type_repr(a) for a in cls.__args__)
+        return f"{name}[{args}]"
+    return cls.__name__
+
+
+def _get_doc(t: mi.Type) -> str:
+    assert hasattr(t, "cls")
+    if hasattr(t.cls, "__origin__"):
+        doc = getattr(t.cls.__origin__, "__doc__")
+    else:
+        doc = getattr(t.cls, "__doc__", "")
+    if not doc:
+        return ""
+    if isinstance(t, mi.EnumType):
+        if doc == "An enumeration.":
+            return ""
+    elif isinstance(t, (mi.NamedTupleType, mi.DataclassType)):
+        if doc.startswith(f"{t.cls.__name__}(") and doc.endswith(")"):
+            return ""
+    return doc
+
+
 def _build_name_map(component_types: dict[Any, mi.Type]) -> dict[Any, str]:
     """A mapping from nameable subcomponents to a generated name.
 
@@ -142,7 +171,7 @@ def _build_name_map(component_types: dict[Any, mi.Type]) -> dict[Any, str]:
     names: dict[str, Any] = {}
 
     for cls in component_types:
-        name = normalize(cls.__name__)
+        name = normalize(_get_class_name(cls))
         if name in names:
             old = names.pop(name)
             conflicts.add(name)
@@ -154,20 +183,6 @@ def _build_name_map(component_types: dict[Any, mi.Type]) -> dict[Any, str]:
     return {v: k for k, v in names.items()}
 
 
-def _has_nondefault_docstring(t: mi.Type) -> bool:
-    """Check if a type has a user-defined docstring.
-
-    Some types like Enum or Dataclass generate a default docstring."""
-    if not (doc := getattr(t.cls, "__doc__", None)):  # type: ignore
-        return False
-
-    if isinstance(t, mi.EnumType):
-        return doc != "An enumeration."
-    elif isinstance(t, (mi.NamedTupleType, mi.DataclassType)):
-        return not (doc.startswith(f"{t.cls.__name__}(") and doc.endswith(")"))
-    return True
-
-
 def _to_schema(
     t: mi.Type, name_map: dict[Any, str], ref_template: str, check_ref: bool = True
 ) -> dict[str, Any]:
@@ -302,13 +317,13 @@ def _to_schema(
         schema["enum"] = sorted(t.values)
     elif isinstance(t, mi.EnumType):
         schema.setdefault("title", t.cls.__name__)
-        if _has_nondefault_docstring(t):
-            schema.setdefault("description", t.cls.__doc__)
+        if doc := _get_doc(t):
+            schema.setdefault("description", doc)
         schema["enum"] = sorted(e.value for e in t.cls)
     elif isinstance(t, mi.StructType):
-        schema.setdefault("title", t.cls.__name__)
-        if _has_nondefault_docstring(t):
-            schema.setdefault("description", t.cls.__doc__)
+        schema.setdefault("title", _get_class_name(t.cls))
+        if doc := _get_doc(t):
+            schema.setdefault("description", doc)
         required = []
         names = []
         fields = []
@@ -347,8 +362,8 @@ def _to_schema(
                 schema["additionalProperties"] = False
     elif isinstance(t, (mi.TypedDictType, mi.DataclassType, mi.NamedTupleType)):
         schema.setdefault("title", t.cls.__name__)
-        if _has_nondefault_docstring(t):
-            schema.setdefault("description", t.cls.__doc__)
+        if doc := _get_doc(t):
+            schema.setdefault("description", doc)
         names = []
         fields = []
         required = []
diff --git a/msgspec/inspect.py b/msgspec/inspect.py
index 64f36e9..12ce579 100644
--- a/msgspec/inspect.py
+++ b/msgspec/inspect.py
@@ -5,7 +5,15 @@ import decimal
 import enum
 import uuid
 from collections.abc import Iterable
-from typing import Any, Final, Literal, Tuple, Type as typing_Type, Union
+from typing import (
+    Any,
+    Final,
+    Literal,
+    Tuple,
+    Type as typing_Type,
+    TypeVar,
+    Union,
+)
 
 try:
     from types import UnionType as _types_UnionType  # type: ignore
@@ -22,8 +30,8 @@ from ._core import (  # type: ignore
 from ._utils import (  # type: ignore
     _CONCRETE_TYPES,
     _AnnotatedAlias,
+    get_class_annotations as _get_class_annotations,
     get_dataclass_info as _get_dataclass_info,
-    get_type_hints as _get_type_hints,
 )
 
 __all__ = (
@@ -691,12 +699,12 @@ class _Translator:
         self.type_hints = {}
         self.cache = {}
 
-    def _get_type_hints(self, t):
-        """A cached version of `get_type_hints`"""
+    def _get_class_annotations(self, t):
+        """A cached version of `get_class_annotations`"""
         try:
             return self.type_hints[t]
         except KeyError:
-            out = self.type_hints[t] = _get_type_hints(t)
+            out = self.type_hints[t] = _get_class_annotations(t)
             return out
 
     def run(self):
@@ -763,6 +771,10 @@ class _Translator:
     ):
         if t is Any:
             return AnyType()
+        elif isinstance(t, TypeVar):
+            if t.__bound__ is not None:
+                return self.translate(t.__bound__)
+            return AnyType()
         elif t is None or t is type(None):
             return NoneType()
         elif t is bool:
@@ -844,11 +856,12 @@ class _Translator:
         elif _is_enum(t):
             return EnumType(t)
         elif _is_struct(t):
-            if t in self.cache:
-                return self.cache[t]
+            cls = t[args] if args else t
+            if cls in self.cache:
+                return self.cache[cls]
             config = t.__struct_config__
-            self.cache[t] = out = StructType(
-                t,
+            self.cache[cls] = out = StructType(
+                cls,
                 (),
                 tag_field=config.tag_field,
                 tag=config.tag,
@@ -856,7 +869,7 @@ class _Translator:
                 forbid_unknown_fields=config.forbid_unknown_fields,
             )
 
-            hints = self._get_type_hints(t)
+            hints = self._get_class_annotations(cls)
             npos = len(t.__struct_fields__) - len(t.__struct_defaults__)
             fields = []
             for name, encode_name, default_obj in zip(
@@ -892,7 +905,7 @@ class _Translator:
             if t in self.cache:
                 return self.cache[t]
             self.cache[t] = out = TypedDictType(t, ())
-            hints = self._get_type_hints(t)
+            hints = self._get_class_annotations(t)
             if hasattr(t, "__required_keys__"):
                 required = set(t.__required_keys__)
             elif t.__total__:
@@ -945,7 +958,7 @@ class _Translator:
             if t in self.cache:
                 return self.cache[t]
             self.cache[t] = out = NamedTupleType(t, ())
-            hints = self._get_type_hints(t)
+            hints = self._get_class_annotations(t)
             out.fields = tuple(
                 Field(
                     name=name,


Test patch:
diff --git a/tests/test_inspect.py b/tests/test_inspect.py
index 489b1cb..2fc9cda 100644
--- a/tests/test_inspect.py
+++ b/tests/test_inspect.py
@@ -14,6 +14,7 @@ from typing import (
     Dict,
     Final,
     FrozenSet,
+    Generic,
     List,
     Literal,
     NamedTuple,
@@ -21,6 +22,7 @@ from typing import (
     Set,
     Tuple,
     TypedDict,
+    TypeVar,
     Union,
 )
 
@@ -42,6 +44,8 @@ except ImportError:
 
 PY39 = sys.version_info[:2] >= (3, 9)
 
+T = TypeVar("T")
+
 
 def type_index(typ, args):
     try:
@@ -81,6 +85,15 @@ def test_any():
     assert mi.type_info(Any) == mi.AnyType()
 
 
+def test_typevar():
+    assert mi.type_info(T) == mi.AnyType()
+
+
+def test_bound_typevar():
+    T = TypeVar("T", bound=Union[int, str])
+    assert mi.type_info(T) == mi.UnionType((mi.IntType(), mi.StrType()))
+
+
 def test_none():
     assert mi.type_info(None) == mi.NoneType()
 
@@ -428,6 +441,30 @@ def test_struct_encode_name():
     assert mi.type_info(Example) == sol
 
 
+def test_generic_struct():
+    class Example(msgspec.Struct, Generic[T]):
+        a: T
+        b: List[T]
+
+    sol = mi.StructType(
+        Example,
+        fields=(
+            mi.Field("a", "a", mi.AnyType()),
+            mi.Field("b", "b", mi.ListType(mi.AnyType())),
+        ),
+    )
+    assert mi.type_info(Example) == sol
+
+    sol = mi.StructType(
+        Example[int],
+        fields=(
+            mi.Field("a", "a", mi.IntType()),
+            mi.Field("b", "b", mi.ListType(mi.IntType())),
+        ),
+    )
+    assert mi.type_info(Example[int]) == sol
+
+
 def test_typing_namedtuple():
     class Example(NamedTuple):
         a: str
diff --git a/tests/test_schema.py b/tests/test_schema.py
index db01bfd..f08c446 100644
--- a/tests/test_schema.py
+++ b/tests/test_schema.py
@@ -10,6 +10,7 @@ from typing import (
     Any,
     Dict,
     FrozenSet,
+    Generic,
     List,
     Literal,
     NamedTuple,
@@ -17,6 +18,7 @@ from typing import (
     Set,
     Tuple,
     TypedDict,
+    TypeVar,
     Union,
 )
 
@@ -35,6 +37,9 @@ except ImportError:
         pytestmark = pytest.mark.skip("Annotated types not available")
 
 
+T = TypeVar("T")
+
+
 def type_index(typ, args):
     try:
         return typ[args]
@@ -863,6 +868,91 @@ def test_struct_unset_fields():
     }
 
 
+def test_generic_struct():
+    class Ex(msgspec.Struct, Generic[T]):
+        """An example docstring"""
+
+        x: T
+        y: List[T]
+
+    assert msgspec.json.schema(Ex) == {
+        "$ref": "#/$defs/Ex",
+        "$defs": {
+            "Ex": {
+                "title": "Ex",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "x": {},
+                    "y": {"type": "array"},
+                },
+                "required": ["x", "y"],
+            },
+        },
+    }
+
+    assert msgspec.json.schema(Ex[int]) == {
+        "$ref": "#/$defs/Ex_int_",
+        "$defs": {
+            "Ex_int_": {
+                "title": "Ex[int]",
+                "description": "An example docstring",
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer"},
+                    "y": {"type": "array", "items": {"type": "integer"}},
+                },
+                "required": ["x", "y"],
+            },
+        },
+    }
+
+
+def test_generic_struct_tagged_union():
+    class Point(msgspec.Struct, Generic[T], tag=True):
+        x: T
+        y: T
+
+    class Point3D(Point[T]):
+        z: T
+
+    sol = {
+        "anyOf": [{"$ref": "#/$defs/Point_int_"}, {"$ref": "#/$defs/Point3D_int_"}],
+        "discriminator": {
+            "mapping": {
+                "Point": "#/$defs/Point_int_",
+                "Point3D": "#/$defs/Point3D_int_",
+            },
+            "propertyName": "type",
+        },
+        "$defs": {
+            "Point_int_": {
+                "properties": {
+                    "type": {"enum": ["Point"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                },
+                "required": ["type", "x", "y"],
+                "title": "Point[int]",
+                "type": "object",
+            },
+            "Point3D_int_": {
+                "properties": {
+                    "type": {"enum": ["Point3D"]},
+                    "x": {"type": "integer"},
+                    "y": {"type": "integer"},
+                    "z": {"type": "integer"},
+                },
+                "required": ["type", "x", "y", "z"],
+                "title": "Point3D[int]",
+                "type": "object",
+            },
+        },
+    }
+    res = msgspec.json.schema(Union[Point[int], Point3D[int]])
+    assert res == sol
+
+
 @pytest.mark.parametrize(
     "field, constraint",
     [