Fix handling of type-var syntax and types.GenericAlias

Author: provinzkraut

src/msgspec/_core.c

@@ -504,6 +504,7 @@ typedef struct {
     PyObject *typing_final;
     PyObject *typing_generic;
     PyObject *typing_generic_alias;
+    PyObject *types_generic_alias;
     PyObject *typing_annotated_alias;
     PyObject *concrete_types;
     PyObject *get_type_hints;
@@ -4948,6 +4949,24 @@ is_dataclass_or_attrs_class(TypeNodeCollectState *state, PyObject *t) {
     );
 }
 
+static MS_INLINE PyObject*
+convert_types_generic_alias(TypeNodeCollectState *state, PyObject *obj, PyObject *origin, PyObject *args) {
+    // if 'obj' is a 'types.GenericAlias', convert it into a 'typing._GenericAlias', so
+    // we can cache type info on it. 'types.GenericAlias' has __slots__, so caching on
+    // it directly does not work.
+    // it's unlikely to hit this case, as it will mostly occur when subclassing a
+    // built-in container generic, such as 'collections.abc.Mapping'
+
+    if (MS_UNLIKELY(Py_TYPE(obj) == (PyTypeObject *)state->mod->types_generic_alias)) {
+        PyObject *genericAliasArgsList = Py_BuildValue("OO", origin, args);
+
+        PyObject *newGenericAlias = PyObject_CallObject(state->mod->typing_generic_alias, genericAliasArgsList);
+        Py_DECREF(genericAliasArgsList);
+        return newGenericAlias;
+    }
+    return obj;
+}
+
 static int
 typenode_collect_type(TypeNodeCollectState *state, PyObject *obj) {
     int out = 0;
@@ -5029,7 +5048,7 @@ typenode_collect_type(TypeNodeCollectState *state, PyObject *obj) {
         ms_is_struct_cls(t) ||
         (origin != NULL && ms_is_struct_cls(origin))
     ) {
-        out = typenode_collect_struct(state, t);
+        out = typenode_collect_struct(state, convert_types_generic_alias(state, t, origin, args));
     }
     else if (Py_TYPE(t) == state->mod->EnumMetaType) {
         out = typenode_collect_enum(state, t);
@@ -5137,7 +5156,7 @@ typenode_collect_type(TypeNodeCollectState *state, PyObject *obj) {
         is_dataclass_or_attrs_class(state, t) ||
         (origin != NULL && is_dataclass_or_attrs_class(state, origin))
     ) {
-        out = typenode_collect_dataclass(state, t);
+        out = typenode_collect_dataclass(state, convert_types_generic_alias(state, t, origin, args));
     }
     else {
         if (origin != NULL) {
@@ -22291,6 +22310,7 @@ msgspec_clear(PyObject *m)
     Py_CLEAR(st->typing_final);
     Py_CLEAR(st->typing_generic);
     Py_CLEAR(st->typing_generic_alias);
+    Py_CLEAR(st->types_generic_alias);
     Py_CLEAR(st->typing_annotated_alias);
     Py_CLEAR(st->concrete_types);
     Py_CLEAR(st->get_type_hints);
@@ -22365,6 +22385,7 @@ msgspec_traverse(PyObject *m, visitproc visit, void *arg)
     Py_VISIT(st->typing_final);
     Py_VISIT(st->typing_generic);
     Py_VISIT(st->typing_generic_alias);
+    Py_VISIT(st->types_generic_alias);
     Py_VISIT(st->typing_annotated_alias);
     Py_VISIT(st->concrete_types);
     Py_VISIT(st->get_type_hints);
@@ -22598,6 +22619,7 @@ PyInit__core(void)
     temp_module = PyImport_ImportModule("types");
     if (temp_module == NULL) return NULL;
     SET_REF(types_uniontype, "UnionType");
+    SET_REF(types_generic_alias, "GenericAlias");
     Py_DECREF(temp_module);
 #endif
 

src/msgspec/_utils.py

@@ -1,6 +1,7 @@
 # type: ignore
 import collections
 import sys
+import types
 import typing
 from typing import _AnnotatedAlias  # noqa: F401
 
@@ -22,6 +23,8 @@ def get_type_hints(obj):
     return _get_type_hints(obj, include_extras=True)
 
 
+PY_31PLUS = sys.version_info >= (3, 12)
+
 # The `is_class` argument was new in 3.11, but was backported to 3.9 and 3.10.
 # It's _likely_ to be available for 3.9/3.10, but may not be. Easiest way to
 # check is to try it and see. This check can be removed when we drop support
@@ -110,13 +113,23 @@ def inner(c, scope):
             cls = c
             new_scope = {}
         else:
-            cls = getattr(c, "__origin__", None)
+            cls = typing.get_origin(c)
             if cls in (None, object, typing.Generic) or cls in mapping:
                 return
-            params = cls.__parameters__
-            args = tuple(_apply_params(a, scope) for a in c.__args__)
-            assert len(params) == len(args)
-            mapping[cls] = new_scope = dict(zip(params, args))
+
+            # it's a built-in generic that has unresolved type vars. in this case,
+            # parameters and args are stored on the generic, not the __origin__
+            if isinstance(c, types.GenericAlias) or (
+                isinstance(c, typing._GenericAlias)
+                and not hasattr(cls, "__parameters__")
+            ):
+                new_scope = dict(zip(c.__parameters__, typing.get_args(c)))
+            else:
+                params = cls.__parameters__
+                args = tuple(_apply_params(a, scope) for a in typing.get_args(c))
+                assert len(params) == len(args)
+                new_scope = dict(zip(params, args))
+            mapping[cls] = new_scope
 
         if issubclass(cls, typing.Generic):
             bases = getattr(cls, "__orig_bases__", cls.__bases__)
@@ -154,6 +167,11 @@ def get_class_annotations(obj):
 
         mapping = typevar_mappings.get(cls)
         cls_locals = dict(vars(cls))
+
+        if PY_31PLUS:
+            # resolve type parameters (e.g. class Foo[T]: pass)
+            cls_locals.update({p.__name__: p for p in cls.__type_params__})
+
         cls_globals = getattr(sys.modules.get(cls.__module__, None), "__dict__", {})
 
         ann = _get_class_annotations(cls)

tests/unit/test_common.py

@@ -1425,6 +1425,35 @@ class Ex(Struct, Generic[T], array_like=array_like):
         with pytest.raises(ValidationError, match="Expected `str`, got `int`"):
             proto.decode(msg, type=Ex[str])
 
+    @py312_plus
+    def test_generic_with_typevar_syntax(self, proto):
+        source = """
+        from msgspec import Struct
+        from typing import List
+        class Ex[T](Struct):
+            x: T
+            y: List[T]
+        """
+
+        with temp_module(source) as mod:
+            sol = mod.Ex(1, [1, 2])
+            msg = proto.encode(sol)
+
+            res = proto.decode(msg, type=mod.Ex)
+            assert res == sol
+
+            res = proto.decode(msg, type=mod.Ex[int])
+            assert res == sol
+
+            res = proto.decode(msg, type=mod.Ex[Union[int, str]])
+            assert res == sol
+
+            res = proto.decode(msg, type=mod.Ex[float])
+            assert type(res.x) is float
+
+            with pytest.raises(ValidationError, match="Expected `str`, got `int`"):
+                proto.decode(msg, type=mod.Ex[str])
+
     @pytest.mark.parametrize("array_like", [False, True])
     def test_recursive_generic_struct(self, proto, array_like):
         source = f"""
@@ -1516,6 +1545,95 @@ def test_unbound_typevars_with_constraints_unsupported(self, proto):
 
         assert "Unbound TypeVar `~T` has constraints" in str(rec.value)
 
+    @pytest.mark.parametrize(
+        "future",
+        [pytest.param(False, id="no future"), pytest.param(False, id="future")],
+    )
+    @pytest.mark.parametrize(
+        "mapping_type", ["collections.abc.Mapping", "typing.Mapping"]
+    )
+    def test_inherited_builtin_generic(self, mapping_type: str, future: bool):
+        source = f"""
+            from msgspec import Struct, StructMeta
+            import collections
+            import abc
+            import typing
+
+            T = typing.TypeVar("T")
+
+            class CombinedMeta(StructMeta, abc.ABCMeta):
+                pass
+
+            class Foo({mapping_type}[str, T], Struct, typing.Generic[T], metaclass=CombinedMeta):
+                data: dict[str, T]
+
+                def __getitem__(self, x):
+                    return self.data[x]
+
+                def __len__(self):
+                    return len(self.data)
+
+                def __iter__(self):
+                    return iter(self.data)
+            """
+
+        if future:
+            source = "from __future__ import annotations\n" + source
+
+        with temp_module(source) as mod, pytest.raises(ValidationError):
+            msgspec.msgpack.decode(
+                msgspec.msgpack.encode(mod.Foo({"x": "foo"})), type=mod.Foo[int]
+            )
+
+        msgspec.msgpack.decode(
+            msgspec.msgpack.encode(mod.Foo({"x": 1})), type=mod.Foo[int]
+        )
+
+    @pytest.mark.parametrize(
+        "future",
+        [pytest.param(False, id="no future"), pytest.param(False, id="future")],
+    )
+    @pytest.mark.parametrize(
+        "mapping_type", ["collections.abc.Mapping", "typing.Mapping"]
+    )
+    @py312_plus
+    def test_inherited_builtin_generic_typevar_syntax(
+        self, mapping_type: str, future: bool
+    ):
+        source = f"""
+            from msgspec import Struct, StructMeta
+            import collections
+            import abc
+            import typing
+
+            class CombinedMeta(StructMeta, abc.ABCMeta):
+                pass
+
+            class Foo[T]({mapping_type}[str, T], Struct, metaclass=CombinedMeta):
+                data: dict[str, T]
+
+                def __getitem__(self, x):
+                    return self.data[x]
+
+                def __len__(self):
+                    return len(self.data)
+
+                def __iter__(self):
+                    return iter(self.data)
+            """
+
+        if future:
+            source = "from __future__ import annotations\n" + source
+
+        with temp_module(source) as mod, pytest.raises(ValidationError):
+            msgspec.msgpack.decode(
+                msgspec.msgpack.encode(mod.Foo({"x": "foo"})), type=mod.Foo[int]
+            )
+
+        msgspec.msgpack.decode(
+            msgspec.msgpack.encode(mod.Foo({"x": 1})), type=mod.Foo[int]
+        )
+
 
 class TestStructPostInit:
     @pytest.mark.parametrize("array_like", [False, True])
@@ -1697,6 +1815,39 @@ class Ex(Generic[T]):
             assert "`$.b.a`" in str(rec.value)
             assert "Expected `int`, got `str`" in str(rec.value)
 
+    @pytest.mark.parametrize("module", ["dataclasses", "attrs"])
+    @py312_plus
+    def test_typevar_syntax(self, module, proto):
+        pytest.importorskip(module)
+        if module == "dataclasses":
+            import_ = "from dataclasses import dataclass as decorator"
+        else:
+            import_ = "from attrs import define as decorator"
+
+        source = f"""
+        from __future__ import annotations
+        from typing import Union
+        from msgspec import Struct
+        {import_}
+
+        @decorator
+        class Ex[T]:
+            a: T
+            b: Union[Ex[T], None]
+        """
+
+        with temp_module(source) as mod:
+            msg = mod.Ex(a=1, b=mod.Ex(a=2, b=None))
+            msg2 = mod.Ex(a=1, b=mod.Ex(a="bad", b=None))
+            assert proto.decode(proto.encode(msg), type=mod.Ex) == msg
+            assert proto.decode(proto.encode(msg2), type=mod.Ex) == msg2
+            assert proto.decode(proto.encode(msg), type=mod.Ex[int]) == msg
+
+            with pytest.raises(ValidationError) as rec:
+                proto.decode(proto.encode(msg2), type=mod.Ex[int])
+            assert "`$.b.a`" in str(rec.value)
+            assert "Expected `int`, got `str`" in str(rec.value)
+
     def test_unbound_typevars_use_bound_if_set(self, proto):
         T = TypeVar("T", bound=Union[int, str])
 
@@ -1719,6 +1870,89 @@ def test_unbound_typevars_with_constraints_unsupported(self, proto):
 
         assert "Unbound TypeVar `~T` has constraints" in str(rec.value)
 
+    @pytest.mark.parametrize(
+        "future",
+        [pytest.param(False, id="no future"), pytest.param(False, id="future")],
+    )
+    @pytest.mark.parametrize(
+        "mapping_type", ["collections.abc.Mapping", "typing.Mapping"]
+    )
+    def test_inherited_builtin_generic(self, mapping_type: str, future: bool):
+        source = f"""
+        import typing
+        import dataclasses
+        import collections
+
+        T = typing.TypeVar("T")
+
+        @dataclasses.dataclass
+        class Foo(typing.Generic[T], {mapping_type}[str, T]):
+            data: dict[str, T]
+
+            def __getitem__(self, x):
+                return self.data[x]
+
+            def __len__(self):
+                return len(self.data)
+
+            def __iter__(self):
+                return iter(self.data)
+        """
+
+        if future:
+            source = "from __future__ import annotations\n" + source
+
+        with temp_module(source) as mod, pytest.raises(ValidationError):
+            msgspec.msgpack.decode(
+                msgspec.msgpack.encode(mod.Foo({"x": "foo"})), type=mod.Foo[int]
+            )
+
+        msgspec.msgpack.decode(
+            msgspec.msgpack.encode(mod.Foo({"x": 1})), type=mod.Foo[int]
+        )
+
+    @pytest.mark.parametrize(
+        "future",
+        [pytest.param(False, id="no future"), pytest.param(False, id="future")],
+    )
+    @pytest.mark.parametrize(
+        "mapping_type", ["collections.abc.Mapping", "typing.Mapping"]
+    )
+    @py312_plus
+    def test_inherited_builtin_generic_typevar_syntax(
+        self, mapping_type: str, future: bool
+    ):
+        source = f"""
+        import dataclasses
+        import collections
+        import typing
+
+        @dataclasses.dataclass
+        class Foo[T]({mapping_type}[str, T]):
+            data: dict[str, T]
+
+            def __getitem__(self, x):
+                return self.data[x]
+
+            def __len__(self):
+                return len(self.data)
+
+            def __iter__(self):
+                return iter(self.data)
+        """
+
+        if future:
+            source = "from __future__ import annotations\n" + source
+
+        with temp_module(source) as mod, pytest.raises(ValidationError):
+            msgspec.msgpack.decode(
+                msgspec.msgpack.encode(mod.Foo({"x": "foo"})), type=mod.Foo[int]
+            )
+
+        msgspec.msgpack.decode(
+            msgspec.msgpack.encode(mod.Foo({"x": 1})), type=mod.Foo[int]
+        )
+
 
 class TestStructOmitDefaults:
     def test_omit_defaults(self, proto):