Expose StructMeta in Python API

docs/source/structs.rst

@@ -998,6 +998,51 @@ container types. It is your responsibility to ensure cycles with these objects
 don't occur, as a cycle containing only ``gc=False`` structs will *never* be
 collected (leading to a memory leak).
 
+Struct Metaclasses (Advanced)
+-----------------------------
+
+:class:`msgspec.Struct` is constructed using the :class:`msgspec.StructMeta`
+metaclass_. We can extend :class:`msgspec.StructMeta` to customize this
+construction. For example, we can automatically set ``kw_only=True`` in all
+subclasses of ``KwOnlyStruct``:
+
+.. code-block:: python
+
+   from msgspec import Struct, StructMeta
+
+   class KwOnlyStructMeta(StructMeta):
+       def __new__(mcls, name, bases, namespace, **kwargs):
+           kwargs.setdefault("kw_only", True)
+           return super().__new__(mcls, name, bases, namespace, **kwargs)
+
+   class KwOnlyStruct(Struct, metaclass=KwOnlyStructMeta):
+       pass
+
+which permits the following syntax:
+
+.. code-block:: python
+
+   class A(KwOnlyStruct):
+       a: int = 1
+       b: str = "hello"
+
+   class B(A):
+       c: int
+
+You can also mix :class:`msgspec.StructMeta` with other metaclasses. For
+example, you can create abstract Struct classes with:
+
+.. code-block:: python
+
+   from msgspec import Struct, StructMeta
+   from abc import ABCMeta
+
+   class ABCStructMeta(StructMeta, ABCMeta):
+       pass
+
+   class StructABC(Struct, metaclass=ABCStructMeta):
+       ...
+
 .. _type annotations: https://docs.python.org/3/library/typing.html
 .. _pattern matching: https://docs.python.org/3/reference/compound_stmts.html#the-match-statement
 .. _PEP 636: https://peps.python.org/pep-0636/
@@ -1013,3 +1058,4 @@ collected (leading to a memory leak).
 .. _rich: https://rich.readthedocs.io/en/stable/pretty.html
 .. _keyword-only parameters: https://docs.python.org/3/glossary.html#term-parameter
 .. _lambda: https://docs.python.org/3/tutorial/controlflow.html#lambda-expressions
+.. _metaclass: https://docs.python.org/3/reference/datamodel.html#metaclasses

msgspec/__init__.py

@@ -6,6 +6,7 @@
     MsgspecError,
     Raw,
     Struct,
+    StructMeta,
     UnsetType,
     UNSET,
     NODEFAULT,

msgspec/__init__.pyi

@@ -1,4 +1,5 @@
 import enum
+from inspect import Signature
 from typing import (
     Any,
     Callable,
@@ -20,6 +21,43 @@ from typing_extensions import dataclass_transform, Buffer
 
 from . import inspect, json, msgpack, structs, toml, yaml
 
+class StructMeta(type):
+     __struct_fields__: ClassVar[Tuple[str, ...]]
+     __struct_defaults__: ClassVar[Tuple[Any, ...]]
+     __struct_encode_fields__: ClassVar[Tuple[str, ...]]
+     __match_args__: ClassVar[Tuple[str, ...]]
+     @property
+     def __signature__(self) -> Signature: ...
+     @property
+     def __struct_config__(self) -> structs.StructConfig: ...
+     def __new__(
+         cls,
+         name: str,
+         bases: Tuple[type, ...],
+         namespace: Dict[str, Any],
+         *,
+         tag_field: Optional[str] = None,
+         tag: Union[None, bool, str, int, Callable[[str], Union[str, int]]] = None,
+         rename: Union[
+             None,
+             Literal["lower", "upper", "camel", "pascal", "kebab"],
+             Callable[[str], Optional[str]],
+             Mapping[str, str],
+         ] = None,
+         omit_defaults: bool = False,
+         forbid_unknown_fields: bool = False,
+         frozen: bool = False,
+         eq: bool = True,
+         order: bool = False,
+         kw_only: bool = False,
+         repr_omit_defaults: bool = False,
+         array_like: bool = False,
+         gc: bool = True,
+         weakref: bool = False,
+         dict: bool = False,
+         cache_hash: bool = False,
+     ) -> StructMeta: ...
+
 T = TypeVar("T")
 
 class UnsetType(enum.Enum):
@@ -39,7 +77,7 @@ def field(*, default_factory: Callable[[], T], name: Optional[str] = None) -> T:
 @overload
 def field(*, name: Optional[str] = None) -> Any: ...
 @dataclass_transform(field_specifiers=(field,))
-class Struct:
+class Struct(metaclass=StructMeta):
     __struct_fields__: ClassVar[Tuple[str, ...]]
     __struct_config__: ClassVar[structs.StructConfig]
     __match_args__: ClassVar[Tuple[str, ...]]

msgspec/_core.c

@@ -4943,8 +4943,8 @@ typenode_collect_type(TypeNodeCollectState *state, PyObject *obj) {
         out = typenode_collect_typevar(state, t);
     }
     else if (
-        Py_TYPE(t) == &StructMetaType ||
-        (origin != NULL && Py_TYPE(origin) == &StructMetaType)
+        PyType_IsSubtype(Py_TYPE(t), &StructMetaType) ||
+        (origin != NULL && PyType_IsSubtype(Py_TYPE(origin), &StructMetaType))
     ) {
         out = typenode_collect_struct(state, t);
     }
@@ -5598,7 +5598,7 @@ structmeta_collect_base(StructMetaInfo *info, MsgspecState *mod, PyObject *base)
         return -1;
     }
 
-    if (Py_TYPE(base) != &StructMetaType) {
+    if (!PyType_IsSubtype(Py_TYPE(base), &StructMetaType)) {
         if (((PyTypeObject *)base)->tp_dictoffset) {
             info->has_non_slots_bases = true;
         }
@@ -5805,7 +5805,7 @@ structmeta_process_default(StructMetaInfo *info, PyObject *name) {
         if (default_val == NULL) return -1;
     }
     else if (
-        (Py_TYPE(type) == &StructMetaType) &&
+        (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) &&
         ((StructMetaObject *)type)->frozen != OPT_TRUE
     ) {
         goto error_mutable_struct;
@@ -6800,7 +6800,7 @@ StructInfo_Convert_lock_held(PyObject *obj) {
     PyObject *annotations = NULL;
     StructInfo *info = NULL;
     bool cache_set = false;
-    bool is_struct = Py_TYPE(obj) == &StructMetaType;
+    bool is_struct = PyType_IsSubtype(Py_TYPE(obj), &StructMetaType);
 
     /* Check for a cached StructInfo, and return if one exists */
     if (MS_LIKELY(is_struct)) {
@@ -6818,7 +6818,7 @@ StructInfo_Convert_lock_held(PyObject *obj) {
         }
         PyObject *origin = PyObject_GetAttr(obj, mod->str___origin__);
         if (origin == NULL) return NULL;
-        if (Py_TYPE(origin) != &StructMetaType) {
+        if (!PyType_IsSubtype(Py_TYPE(origin), &StructMetaType)) {
             Py_DECREF(origin);
             PyErr_SetString(
                 PyExc_RuntimeError, "Expected __origin__ to be a Struct type"
@@ -7276,7 +7276,7 @@ static PyTypeObject StructMetaType = {
     .tp_name = "msgspec._core.StructMeta",
     .tp_basicsize = sizeof(StructMetaObject),
     .tp_itemsize = 0,
-    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_TYPE_SUBCLASS | Py_TPFLAGS_HAVE_GC | _Py_TPFLAGS_HAVE_VECTORCALL,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_TYPE_SUBCLASS | Py_TPFLAGS_HAVE_GC | _Py_TPFLAGS_HAVE_VECTORCALL | Py_TPFLAGS_BASETYPE,
     .tp_new = StructMeta_new,
     .tp_dealloc = (destructor) StructMeta_dealloc,
     .tp_clear = (inquiry) StructMeta_clear,
@@ -7915,7 +7915,7 @@ struct_replace(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject
 
     if (!check_positional_nargs(nargs, 1, 1)) return NULL;
     PyObject *obj = args[0];
-    if (Py_TYPE(Py_TYPE(obj)) != &StructMetaType) {
+    if (!PyType_IsSubtype(Py_TYPE(Py_TYPE(obj)), &StructMetaType)) {
         PyErr_SetString(PyExc_TypeError, "`struct` must be a `msgspec.Struct`");
         return NULL;
     }
@@ -7956,7 +7956,7 @@ struct_asdict(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
 {
     if (!check_positional_nargs(nargs, 1, 1)) return NULL;
     PyObject *obj = args[0];
-    if (Py_TYPE(Py_TYPE(obj)) != &StructMetaType) {
+    if (!PyType_IsSubtype(Py_TYPE(Py_TYPE(obj)), &StructMetaType)) {
         PyErr_SetString(PyExc_TypeError, "`struct` must be a `msgspec.Struct`");
         return NULL;
     }
@@ -8014,7 +8014,7 @@ struct_astuple(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
 {
     if (!check_positional_nargs(nargs, 1, 1)) return NULL;
     PyObject *obj = args[0];
-    if (Py_TYPE(Py_TYPE(obj)) != &StructMetaType) {
+    if (!PyType_IsSubtype(Py_TYPE(Py_TYPE(obj)), &StructMetaType)) {
         PyErr_SetString(PyExc_TypeError, "`struct` must be a `msgspec.Struct`");
         return NULL;
     }
@@ -8066,7 +8066,7 @@ struct_force_setattr(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
     PyObject *obj = args[0];
     PyObject *name = args[1];
     PyObject *value = args[2];
-    if (Py_TYPE(Py_TYPE(obj)) != &StructMetaType) {
+    if (!PyType_IsSubtype(Py_TYPE(Py_TYPE(obj)), &StructMetaType)) {
         PyErr_SetString(PyExc_TypeError, "`struct` must be a `msgspec.Struct`");
         return NULL;
     }
@@ -13245,7 +13245,7 @@ mpack_encode_uncommon(EncoderState *self, PyTypeObject *type, PyObject *obj)
     else if (type == &PyBool_Type) {
         return mpack_encode_bool(self, obj);
     }
-    else if (Py_TYPE(type) == &StructMetaType) {
+    else if (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) {
         return mpack_encode_struct(self, obj);
     }
     else if (type == &PyBytes_Type) {
@@ -14352,7 +14352,7 @@ json_encode_uncommon(EncoderState *self, PyTypeObject *type, PyObject *obj) {
     else if (obj == Py_False) {
         return ms_write(self, "false", 5);
     }
-    else if (Py_TYPE(type) == &StructMetaType) {
+    else if (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) {
         return json_encode_struct(self, obj);
     }
     else if (PyTuple_Check(obj)) {
@@ -16481,7 +16481,7 @@ msgspec_msgpack_decode(PyObject *self, PyObject *const *args, Py_ssize_t nargs,
     if (type == NULL || type == mod->typing_any) {
         state.type = &typenode_any;
     }
-    else if (Py_TYPE(type) == &StructMetaType) {
+    else if (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) {
         PyObject *info = StructInfo_Convert(type);
         if (info == NULL) return NULL;
         bool array_like = ((StructMetaObject *)type)->array_like == OPT_TRUE;
@@ -19516,7 +19516,7 @@ msgspec_json_decode(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyO
     if (type == NULL || type == mod->typing_any) {
         state.type = &typenode_any;
     }
-    else if (Py_TYPE(type) == &StructMetaType) {
+    else if (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) {
         PyObject *info = StructInfo_Convert(type);
         if (info == NULL) return NULL;
         bool array_like = ((StructMetaObject *)type)->array_like == OPT_TRUE;
@@ -20084,7 +20084,7 @@ to_builtins(ToBuiltinsState *self, PyObject *obj, bool is_key) {
     else if (PyDict_Check(obj)) {
         return to_builtins_dict(self, obj);
     }
-    else if (Py_TYPE(type) == &StructMetaType) {
+    else if (PyType_IsSubtype(Py_TYPE(type), &StructMetaType)) {
         return to_builtins_struct(self, obj, is_key);
     }
     else if (Py_TYPE(type) == self->mod->EnumMetaType) {
@@ -21983,7 +21983,7 @@ msgspec_convert(PyObject *self, PyObject *args, PyObject *kwargs)
     state.dec_hook = dec_hook;
 
     /* Avoid allocating a new TypeNode for struct types */
-    if (Py_TYPE(pytype) == &StructMetaType) {
+    if (PyType_IsSubtype(Py_TYPE(pytype), &StructMetaType)) {
         PyObject *info = StructInfo_Convert(pytype);
         if (info == NULL) return NULL;
         bool array_like = ((StructMetaObject *)pytype)->array_like == OPT_TRUE;
@@ -22305,7 +22305,13 @@ PyInit__core(void)
         return NULL;
     Py_INCREF(&Unset_Type);
     if (PyModule_AddObject(m, "UnsetType", (PyObject *)&Unset_Type) < 0)
+        return NULL;    
+    Py_INCREF((PyObject *)&StructMetaType);
+    if (PyModule_AddObject(m, "StructMeta", (PyObject *)&StructMetaType) < 0) {
+        Py_DECREF((PyObject *)&StructMetaType);
+        Py_DECREF(m);
         return NULL;
+    }
 
     st = msgspec_get_state(m);