Skip to content

Commit

Permalink
Support typing.Final annotations
Browse files Browse the repository at this point in the history 
This adds support for `typing.Final` annotations. `Final` can be used to
wrap an existing field annotation, marking it as a field that can't be
modified once it's initialized. This has the same semantics as
`frozen=True`, but only for a single field, and not enforced at runtime.

The `Frozen` annotation may be used to wrap field annotations on any
object-like type we support (`msgspec.Struct`, `attrs`, or
`dataclasses`).
  • Loading branch information
@jcrist
jcrist committed Mar 25, 2023
1 parent 4f064c8 commit 73f9041
Show file tree
Hide file tree
Showing 7 changed files with 196 additions and 135 deletions.
1 change: 1 addition & 0 deletions docs/source/supported-types.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ Most combinations of the following types are supported (with a few restrictions)
- `typing.Union`
- `typing.Literal`
- `typing.NewType`
- `typing.Final`
- `typing.NamedTuple` / `collections.namedtuple`
- `typing.TypedDict`

Expand Down
156 changes: 87 additions & 69 deletions msgspec/_core.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -372,6 +372,7 @@ typedef struct {
PyObject *typing_any;
PyObject *typing_literal;
PyObject *typing_classvar;
PyObject *typing_final;
PyObject *typing_generic_alias;
PyObject *typing_annotated_alias;
PyObject *concrete_types;
Expand Down Expand Up @@ -4123,44 +4124,92 @@ typenode_origin_args_metadata(
PyObject *t = obj;
Py_INCREF(t);

/* First strip out meta "wrapper" types (Annotated, NewType) */
/* First strip out meta "wrapper" types (Annotated, NewType, Final) */
while (true) {
if (Py_TYPE(t) == (PyTypeObject *)(state->mod->typing_annotated_alias)) {
/* Handle Annotated */
PyObject *origin = PyObject_GetAttr(t, state->mod->str___origin__);
if (origin == NULL) {
Py_CLEAR(t);
return NULL;
}
assert(t != NULL && origin == NULL && args == NULL);

PyObject *metadata = PyObject_GetAttr(t, state->mod->str___metadata__);
if (metadata == NULL) {
Py_DECREF(origin);
/* Before inspecting attributes, try looking up the object in the
* abstract -> concrete mapping. If present, this is an unparametrized
* collection of some form. This helps avoid compatibility issues in
* Python 3.8, where unparametrized collections still have __args__. */
origin = PyDict_GetItem(state->mod->concrete_types, t);
if (origin != NULL) {
Py_INCREF(origin);
break;
}

/* If `t` is a type instance, no need to inspect further */
if (PyType_CheckExact(t)) {
/* t is a concrete type object. */
break;
}

origin = PyObject_GetAttr(t, state->mod->str___origin__);
if (origin != NULL) {
if (Py_TYPE(t) == (PyTypeObject *)(state->mod->typing_annotated_alias)) {
/* Handle typing.Annotated[...] */
PyObject *metadata = PyObject_GetAttr(t, state->mod->str___metadata__);
if (metadata == NULL) goto error;
for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(metadata); i++) {
PyObject *annot = PyTuple_GET_ITEM(metadata, i);
if (Py_TYPE(annot) == &Meta_Type) {
if (constraints_update(constraints, (Meta *)annot, obj) < 0) {
Py_DECREF(metadata);
goto error;
}
}
}
Py_DECREF(metadata);
Py_DECREF(t);
return NULL;
t = origin;
origin = NULL;
continue;
}

for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(metadata); i++) {
PyObject *annot = PyTuple_GET_ITEM(metadata, i);
if (Py_TYPE(annot) == &Meta_Type) {
if (constraints_update(constraints, (Meta *)annot, obj) < 0) {
Py_DECREF(metadata);
Py_DECREF(origin);
else {
args = PyObject_GetAttr(t, state->mod->str___args__);
if (args != NULL) {
if (!PyTuple_Check(args)) {
PyErr_SetString(PyExc_TypeError, "__args__ must be a tuple");
goto error;
}
if (origin == state->mod->typing_final) {
/* Handle typing.Final[...] */
PyObject *temp = PyTuple_GetItem(args, 0);
if (temp == NULL) goto error;
Py_CLEAR(args);
Py_CLEAR(origin);
Py_DECREF(t);
return NULL;
Py_INCREF(temp);
t = temp;
continue;
}
}
else {
/* Custom non-parametrized generics won't have __args__
* set. Ignore __args__ error */
PyErr_Clear();
}
/* Lookup __origin__ in the mapping, in case it's a supported
* abstract type. Equal to `origin = mapping.get(origin, origin)` */
PyObject *temp = PyDict_GetItem(state->mod->concrete_types, origin);
if (temp != NULL) {
Py_DECREF(origin);
Py_INCREF(temp);
origin = temp;
}
break;
}
Py_DECREF(metadata);
Py_DECREF(t);
t = origin;
}
else {
/* Handle NewType */
PyErr_Clear();

/* Check for NewType */
PyObject *supertype = PyObject_GetAttr(t, state->mod->str___supertype__);
if (supertype != NULL) {
/* It's a newtype, use the wrapped type and loop again */
Py_DECREF(t);
t = supertype;
continue;
}
else {
PyErr_Clear();
Expand All @@ -4169,59 +4218,25 @@ typenode_origin_args_metadata(
}
}

/* At this point `t` is a concrete type. Next check for generic types,
* extracting `__origin__` and `__args__`. This lets us normalize how
* we check for collection types later */
if ((origin = PyDict_GetItem(state->mod->concrete_types, t)) != NULL) {
Py_INCREF(origin);
}
#if PY_VERSION_HEX >= 0x030a00f0
else if (Py_TYPE(t) == (PyTypeObject *)(state->mod->types_uniontype)) {
if (Py_TYPE(t) == (PyTypeObject *)(state->mod->types_uniontype)) {
/* Handle types.UnionType unions (`int | float | ...`) */
args = PyObject_GetAttr(t, state->mod->str___args__);
if (args == NULL) {
Py_DECREF(t);
return NULL;
}
if (args == NULL) goto error;
origin = state->mod->typing_union;
Py_INCREF(origin);
}
#endif
else {
origin = PyObject_GetAttr(t, state->mod->str___origin__);
if (origin == NULL) {
/* Not a generic */
PyErr_Clear();
}
else {
/* Lookup __origin__ in the mapping, in case it's a supported
* abstract type */
PyObject *temp = PyDict_GetItem(state->mod->concrete_types, origin);
if (temp != NULL) {
Py_DECREF(origin);
Py_INCREF(temp);
origin = temp;
}
args = PyObject_GetAttr(t, state->mod->str___args__);
if (args == NULL) {
/* Custom non-parametrized generics won't have __args__ set.
* Ignore __args__ error */
PyErr_Clear();
}
else {
if (!PyTuple_Check(args)) {
PyErr_SetString(PyExc_TypeError, "__args__ must be a tuple");
Py_DECREF(t);
Py_DECREF(origin);
Py_DECREF(args);
return NULL;
}
}
}
}

*out_origin = origin;
*out_args = args;
return t;

error:
Py_XDECREF(t);
Py_XDECREF(origin);
Py_XDECREF(args);
return NULL;
}

static int
Expand Down Expand Up @@ -10443,7 +10458,7 @@ mpack_encode_struct(EncoderState *self, PyObject *obj)
actual_len--;
}
else {
if (mpack_encode_str(self, key) < 0) goto cleanup;
if (mpack_encode_str(self, key) < 0) goto cleanup;
if (mpack_encode(self, val) < 0) goto cleanup;
}
}
Expand All @@ -10458,7 +10473,7 @@ mpack_encode_struct(EncoderState *self, PyObject *obj)
actual_len--;
}
else {
if (mpack_encode_str(self, key) < 0) goto cleanup;
if (mpack_encode_str(self, key) < 0) goto cleanup;
if (mpack_encode(self, val) < 0) goto cleanup;
}
}
Expand Down Expand Up @@ -18606,6 +18621,7 @@ msgspec_clear(PyObject *m)
Py_CLEAR(st->typing_any);
Py_CLEAR(st->typing_literal);
Py_CLEAR(st->typing_classvar);
Py_CLEAR(st->typing_final);
Py_CLEAR(st->typing_generic_alias);
Py_CLEAR(st->typing_annotated_alias);
Py_CLEAR(st->concrete_types);
Expand Down Expand Up @@ -18686,6 +18702,7 @@ msgspec_traverse(PyObject *m, visitproc visit, void *arg)
Py_VISIT(st->typing_any);
Py_VISIT(st->typing_literal);
Py_VISIT(st->typing_classvar);
Py_VISIT(st->typing_final);
Py_VISIT(st->typing_generic_alias);
Py_VISIT(st->typing_annotated_alias);
Py_VISIT(st->concrete_types);
Expand Down Expand Up @@ -18893,6 +18910,7 @@ PyInit__core(void)
SET_REF(typing_any, "Any");
SET_REF(typing_literal, "Literal");
SET_REF(typing_classvar, "ClassVar");
SET_REF(typing_final, "Final");
SET_REF(typing_generic_alias, "_GenericAlias");
Py_DECREF(temp_module);

Expand Down
68 changes: 26 additions & 42 deletions msgspec/_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,50 +34,34 @@ def get_type_hints(obj):


# A mapping from a type annotation (or annotation __origin__) to the concrete
# python type that msgspec will use when decoding. Note that non-collection
# types don't strict need to be in this mapping. Common ones are added to avoid
# an unnecessary `getattr(t, "__origin__", None)` call on them.
# THIS IS PRIVATE FOR A REASON. DON'T MUCK WITH THIS.
# python type that msgspec will use when decoding. THIS IS PRIVATE FOR A
# REASON. DON'T MUCK WITH THIS.
_CONCRETE_TYPES = {
t: t
for t in [
None,
bool,
int,
float,
str,
bytes,
bytearray,
list,
tuple,
set,
frozenset,
dict,
]
list: list,
tuple: tuple,
set: set,
frozenset: frozenset,
dict: dict,
typing.List: list,
typing.Tuple: tuple,
typing.Set: set,
typing.FrozenSet: frozenset,
typing.Dict: dict,
typing.Collection: list,
typing.MutableSequence: list,
typing.Sequence: list,
typing.MutableMapping: dict,
typing.Mapping: dict,
typing.MutableSet: set,
typing.AbstractSet: set,
collections.abc.Collection: list,
collections.abc.MutableSequence: list,
collections.abc.Sequence: list,
collections.abc.MutableSet: set,
collections.abc.Set: set,
collections.abc.MutableMapping: dict,
collections.abc.Mapping: dict,
}
_CONCRETE_TYPES.update(
{
typing.List: list,
typing.Tuple: tuple,
typing.Set: set,
typing.FrozenSet: frozenset,
typing.Dict: dict,
typing.Collection: list,
typing.MutableSequence: list,
typing.Sequence: list,
typing.MutableMapping: dict,
typing.Mapping: dict,
typing.MutableSet: set,
typing.AbstractSet: set,
collections.abc.Collection: list,
collections.abc.MutableSequence: list,
collections.abc.Sequence: list,
collections.abc.MutableSet: set,
collections.abc.Set: set,
collections.abc.MutableMapping: dict,
collections.abc.Mapping: dict,
}
)


def get_typeddict_hints(obj):
Expand Down
50 changes: 27 additions & 23 deletions msgspec/inspect.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import enum
import uuid
from collections.abc import Iterable
from typing import Any, Literal, Tuple, Type as typing_Type, Union
from typing import Any, Final, Literal, Tuple, Type as typing_Type, Union

try:
from types import UnionType as _types_UnionType
Expand Down Expand Up @@ -611,34 +611,38 @@ def type_info(type: Any, *, protocol: Literal[None, "msgpack", "json"] = None) -

# Implementation details
def _origin_args_metadata(t):
# Strip Annotated and NewType wrappers until we hit a concrete base type
# Strip wrappers (Annotated, NewType, Final) until we hit a concrete type
metadata = []
while True:
supertype = getattr(t, "__supertype__", None)
if supertype is not None:
t = supertype
elif type(t) is _AnnotatedAlias:
metadata.extend(m for m in t.__metadata__ if type(m) is msgspec.Meta)
t = t.__origin__
else:
origin = _CONCRETE_TYPES.get(t)
if origin is not None:
args = None
break

if type(t) is _types_UnionType:
args = t.__args__
t = Union
else:
try:
t = _CONCRETE_TYPES[t]
args = None
except Exception:
try:
origin = t.__origin__
except AttributeError:
args = None
origin = getattr(t, "__origin__", None)
if origin is not None:
if type(t) is _AnnotatedAlias:
metadata.extend(m for m in t.__metadata__ if type(m) is msgspec.Meta)
t = origin
elif origin == Final:
t = t.__args__[0]
else:
args = getattr(t, "__args__", None)
t = _CONCRETE_TYPES.get(origin, origin)
return t, args, tuple(metadata)
origin = _CONCRETE_TYPES.get(origin, origin)
break
else:
supertype = getattr(t, "__supertype__", None)
if supertype is not None:
t = supertype
else:
origin = t
args = None
break

if type(origin) is _types_UnionType:
args = origin.__args__
origin = Union
return origin, args, tuple(metadata)


def _is_struct(t):
Expand Down
Loading

0 comments on commit 73f9041

Please sign in to comment.