From f18370f26679574309e8647ccb65a7d923fb289b Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 00:05:57 -0500 Subject: [PATCH 1/6] Add initial support for struct class-level options Adds support for struct class-level options. These can be set as parameters to the metaclass. For example: ```python import msgspec class Point(msgspec.Struct, immutable=True): x: float y: float ``` Currently these options don't do anything (TODO), but the following options are defined: - `immutable`: makes a struct object immutable, and adds a `__hash__` implementation. - `asarray`: whether to serialize the struct as an array instead of as a dict. --- msgspec/core.c | 59 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 56 insertions(+), 3 deletions(-) diff --git a/msgspec/core.c b/msgspec/core.c index 3d2a6136..ed99a1cd 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -713,6 +713,8 @@ typedef struct { PyObject *struct_defaults; Py_ssize_t *struct_offsets; TypeNode **struct_types; + char immutable; + char asarray; } StructMetaObject; static PyTypeObject StructMixinType; @@ -722,6 +724,12 @@ static PyTypeObject StructMixinType; #define StructMeta_GET_DEFAULTS(s) (((StructMetaObject *)(s))->struct_defaults); #define StructMeta_GET_OFFSETS(s) (((StructMetaObject *)(s))->struct_offsets); +#define OPT_UNSET -1 +#define OPT_FALSE 0 +#define OPT_TRUE 1 +#define STRUCT_MERGE_OPTIONS(opt1, opt2) (((opt2) != OPT_UNSET) ? (opt2) : (opt1)) + + static Py_ssize_t StructMeta_get_field_index(StructMetaObject *self, char * key, Py_ssize_t key_size, Py_ssize_t *pos) { const char *field; @@ -764,10 +772,15 @@ StructMeta_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) PyObject *default_val, *field; Py_ssize_t nfields, ndefaults, i, j, k; Py_ssize_t *offsets = NULL, *base_offsets; + int arg_immutable = -1, arg_asarray = -1, immutable = -1, asarray = -1; + + static char *kwlist[] = {"name", "bases", "dict", "immutable", "asarray", NULL}; /* Parse arguments: (name, bases, dict) */ - if (!PyArg_ParseTuple(args, "UO!O!:StructMeta.__new__", &name, &PyTuple_Type, - &bases, &PyDict_Type, &orig_dict)) + if (!PyArg_ParseTupleAndKeywords( + args, kwargs, "UO!O!|$pp:StructMeta.__new__", kwlist, + &name, &PyTuple_Type, &bases, &PyDict_Type, &orig_dict, + &arg_immutable, &arg_asarray)) return NULL; if (PyDict_GetItemString(orig_dict, "__init__") != NULL) { @@ -806,6 +819,8 @@ StructMeta_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) ); goto error; } + immutable = STRUCT_MERGE_OPTIONS(immutable, ((StructMetaObject *)base)->immutable); + asarray = STRUCT_MERGE_OPTIONS(asarray, ((StructMetaObject *)base)->asarray); base_fields = StructMeta_GET_FIELDS(base); base_defaults = StructMeta_GET_DEFAULTS(base); base_offsets = StructMeta_GET_OFFSETS(base); @@ -834,6 +849,8 @@ StructMeta_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) Py_DECREF(offset); } } + immutable = STRUCT_MERGE_OPTIONS(immutable, arg_immutable); + asarray = STRUCT_MERGE_OPTIONS(asarray, arg_asarray); new_dict = PyDict_Copy(orig_dict); if (new_dict == NULL) @@ -924,7 +941,7 @@ StructMeta_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) if (new_args == NULL) goto error; - cls = (StructMetaObject *) PyType_Type.tp_new(type, new_args, kwargs); + cls = (StructMetaObject *) PyType_Type.tp_new(type, new_args, NULL); if (cls == NULL) goto error; ((PyTypeObject *)cls)->tp_vectorcall = (vectorcallfunc)Struct_vectorcall; @@ -955,6 +972,8 @@ StructMeta_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) cls->struct_fields = fields; cls->struct_defaults = defaults; cls->struct_offsets = offsets; + cls->immutable = immutable; + cls->asarray = asarray; return (PyObject *) cls; error: Py_XDECREF(arg_fields); @@ -1060,6 +1079,20 @@ StructMeta_dealloc(StructMetaObject *self) PyType_Type.tp_dealloc((PyObject *)self); } +static PyObject* +StructMeta_immutable(StructMetaObject *self, void *closure) +{ + if (self->immutable == OPT_TRUE) { Py_RETURN_TRUE; } + else { Py_RETURN_FALSE; } +} + +static PyObject* +StructMeta_asarray(StructMetaObject *self, void *closure) +{ + if (self->asarray == OPT_TRUE) { Py_RETURN_TRUE; } + else { Py_RETURN_FALSE; } +} + static PyObject* StructMeta_signature(StructMetaObject *self, void *closure) { @@ -1159,6 +1192,8 @@ static PyMemberDef StructMeta_members[] = { static PyGetSetDef StructMeta_getset[] = { {"__signature__", (getter) StructMeta_signature, NULL, NULL, NULL}, + {"immutable", (getter) StructMeta_immutable, NULL, NULL, NULL}, + {"asarray", (getter) StructMeta_asarray, NULL, NULL, NULL}, {NULL}, }; @@ -1563,6 +1598,14 @@ PyDoc_STRVAR(Struct__doc__, "Note that mutable default values are deepcopied in the constructor to\n" "prevent accidental sharing.\n" "\n" +"Additional class options can be enabled by passing keywords to the class\n" +"definition (see example below). The following options exist:\n" +"\n" +"- ``immutable``: whether instances of the class are immutable. If true,\n" +" attribute assignment is disabled and a corresponding ``__hash__`` is defined.\n" +"- ``asarray``: whether instances of the class should be serialized as\n" +" MessagePack arrays, rather than dicts (the default).\n" +"\n" "Structs automatically define ``__init__``, ``__eq__``, ``__repr__``, and\n" "``__copy__`` methods. Additional methods can be defined on the class as\n" "needed. Note that ``__init__``/``__new__`` cannot be overridden, but other\n" @@ -1581,6 +1624,16 @@ PyDoc_STRVAR(Struct__doc__, "...\n" ">>> Dog('snickers', breed='corgi')\n" "Dog(name='snickers', breed='corgi', is_good_boy=True)\n" +"\n" +"Additional struct options can be set as part of the class definition. Here\n" +"we define a new `Struct` type for an immutable `Point` object.\n" +"\n" +">>> class Point(Struct, immutable=True):\n" +"... x: float\n" +"... y: float\n" +"...\n" +">>> {Point(1.5, 2.0): 1} # immutable structs are hashable\n" +"{Point(1.5, 2.0): 1}" ); /************************************************************************* From 87d66139552a4440841b7bddfa182c2735a3884c Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 01:16:56 -0500 Subject: [PATCH 2/6] Implement `immutable` option --- msgspec/core.c | 47 ++++++++++++++++++++++++++++ tests/test_struct.py | 73 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+) diff --git a/msgspec/core.c b/msgspec/core.c index ed99a1cd..860170b0 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -1392,6 +1392,10 @@ Struct_vectorcall(PyTypeObject *cls, PyObject *const *args, size_t nargsf, PyObj static int Struct_setattro(PyObject *self, PyObject *key, PyObject *value) { + if (((StructMetaObject *)Py_TYPE(self))->immutable == OPT_TRUE) { + PyErr_Format(PyExc_TypeError, "immutable type: '%s'", Py_TYPE(self)->tp_name); + return -1; + } if (PyObject_GenericSetAttr(self, key, value) < 0) return -1; if (value != NULL && OBJ_IS_GC(value) && !IS_TRACKED(self)) @@ -1455,6 +1459,48 @@ Struct_repr(PyObject *self) { return out; } +/* Hash algorithm borrowed from cpython 3.10's hashing algorithm for tuples. + * See https://github.com/python/cpython/blob/4bcef2bb48b3fd82011a89c1c716421b789f1442/Objects/tupleobject.c#L386-L424 + */ +#if SIZEOF_PY_UHASH_T > 4 +#define MP_HASH_XXPRIME_1 ((Py_uhash_t)11400714785074694791ULL) +#define MP_HASH_XXPRIME_2 ((Py_uhash_t)14029467366897019727ULL) +#define MP_HASH_XXPRIME_5 ((Py_uhash_t)2870177450012600261ULL) +#define MP_HASH_XXROTATE(x) ((x << 31) | (x >> 33)) /* Rotate left 31 bits */ +#else +#define MP_HASH_XXPRIME_1 ((Py_uhash_t)2654435761UL) +#define MP_HASH_XXPRIME_2 ((Py_uhash_t)2246822519UL) +#define MP_HASH_XXPRIME_5 ((Py_uhash_t)374761393UL) +#define MP_HASH_XXROTATE(x) ((x << 13) | (x >> 19)) /* Rotate left 13 bits */ +#endif + +static Py_hash_t +Struct_hash(PyObject *self) { + PyObject *val; + Py_ssize_t i, nfields; + Py_uhash_t acc = MP_HASH_XXPRIME_5; + + if (((StructMetaObject *)Py_TYPE(self))->immutable == OPT_UNSET) { + PyErr_Format(PyExc_TypeError, "unhashable type: '%s'", Py_TYPE(self)->tp_name); + return -1; + } + + nfields = StructMeta_GET_NFIELDS(Py_TYPE(self)); + + for (i = 0; i < nfields; i++) { + Py_uhash_t lane; + val = Struct_get_index(self, i); + if (val == NULL) return -1; + lane = PyObject_Hash(val); + if (lane == (Py_uhash_t)-1) return -1; + acc += lane * MP_HASH_XXPRIME_2; + acc = MP_HASH_XXROTATE(acc); + acc *= MP_HASH_XXPRIME_1; + } + acc += nfields ^ (MP_HASH_XXPRIME_5 ^ 3527539UL); + return (acc == (Py_uhash_t)-1) ? 1546275796 : acc; +} + static PyObject * Struct_richcompare(PyObject *self, PyObject *other, int op) { int status; @@ -1586,6 +1632,7 @@ static PyTypeObject StructMixinType = { .tp_setattro = Struct_setattro, .tp_repr = Struct_repr, .tp_richcompare = Struct_richcompare, + .tp_hash = Struct_hash, .tp_methods = Struct_methods, .tp_getset = StructMixin_getset, }; diff --git a/tests/test_struct.py b/tests/test_struct.py index 7d113089..66486b4c 100644 --- a/tests/test_struct.py +++ b/tests/test_struct.py @@ -645,3 +645,76 @@ def test_struct_handles_missing_attributes(): with pytest.raises(AttributeError, match=match): pickle.dumps(t) + + +def test_struct_option_precedence(): + class Default(Struct): + pass + + assert not Default.immutable + + class Immutable(Struct, immutable=True): + pass + + assert Immutable.immutable + + class NotImmutable(Struct, immutable=False): + pass + + assert not NotImmutable.immutable + + class T(Immutable): + pass + + assert T.immutable + + class T(Immutable, immutable=False): + pass + + assert not T.immutable + + class T(Immutable, Default): + pass + + assert T.immutable + + class T(Default, Immutable): + pass + + assert T.immutable + + class T(Default, NotImmutable, Immutable): + pass + + assert not T.immutable + + +class ImmutablePoint(Struct, immutable=True): + x: int + y: int + + +class TestImmutable: + def test_immutable_objects_no_setattr(self): + p = ImmutablePoint(1, 2) + with pytest.raises(TypeError, match="immutable type: 'ImmutablePoint'"): + p.x = 3 + + def test_immutable_objects_hashable(self): + p1 = ImmutablePoint(1, 2) + p2 = ImmutablePoint(1, 2) + p3 = ImmutablePoint(1, 3) + assert hash(p1) == hash(p2) + assert hash(p1) != hash(p3) + assert p1 == p2 + assert p1 != p3 + + def test_immutable_objects_hash_errors_if_field_unhashable(self): + p = ImmutablePoint(1, [2]) + with pytest.raises(TypeError): + hash(p) + + def test_mutable_objects_hash_errors(self): + p = Point(1, 2) + with pytest.raises(TypeError, match="unhashable type"): + hash(p) From c87d5405e9cf90c138d718f4573e27ba94585c43 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 01:28:34 -0500 Subject: [PATCH 3/6] Add `asarray` option for encoding structs as arrays --- msgspec/core.c | 37 ++++++++++++++++++++++++++----------- 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/msgspec/core.c b/msgspec/core.c index 860170b0..d03482e3 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -1480,7 +1480,7 @@ Struct_hash(PyObject *self) { Py_ssize_t i, nfields; Py_uhash_t acc = MP_HASH_XXPRIME_5; - if (((StructMetaObject *)Py_TYPE(self))->immutable == OPT_UNSET) { + if (((StructMetaObject *)Py_TYPE(self))->immutable != OPT_TRUE) { PyErr_Format(PyExc_TypeError, "unhashable type: '%s'", Py_TYPE(self)->tp_name); return -1; } @@ -2362,21 +2362,36 @@ mp_encode_struct(EncoderState *self, PyObject *obj) PyObject *key, *val, *fields; Py_ssize_t i, len; int status = 0; + bool asarray = ((StructMetaObject *)Py_TYPE(obj))->asarray == OPT_TRUE; fields = StructMeta_GET_FIELDS(Py_TYPE(obj)); len = PyTuple_GET_SIZE(fields); - if (mp_encode_map_header(self, len, "structs") < 0) - return -1; - if (len == 0) - return 0; + + status = ( + asarray ? mp_encode_array_header(self, len, "structs") : + mp_encode_map_header(self, len, "structs") + ); + if (status < 0) return -1; + if (len == 0) return 0; if (Py_EnterRecursiveCall(" while serializing an object")) return -1; - for (i = 0; i < len; i++) { - key = PyTuple_GET_ITEM(fields, i); - val = Struct_get_index(obj, i); - if (val == NULL || mp_encode(self, key) < 0 || mp_encode(self, val) < 0) { - status = -1; - break; + if (asarray) { + for (i = 0; i < len; i++) { + val = Struct_get_index(obj, i); + if (val == NULL || mp_encode(self, val) < 0) { + status = -1; + break; + } + } + } + else { + for (i = 0; i < len; i++) { + key = PyTuple_GET_ITEM(fields, i); + val = Struct_get_index(obj, i); + if (val == NULL || mp_encode(self, key) < 0 || mp_encode(self, val) < 0) { + status = -1; + break; + } } } Py_LeaveRecursiveCall(); From a3212c4f1e2f3e4f33aadf24aa5fe236f75ec5f3 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 08:56:37 -0500 Subject: [PATCH 4/6] Support decoding structs from arrays --- msgspec/core.c | 133 +++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 113 insertions(+), 20 deletions(-) diff --git a/msgspec/core.c b/msgspec/core.c index d03482e3..045f8ccc 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -3983,27 +3983,23 @@ mp_decode_type_ext(DecoderState *self, char op, TypeNode *ctx, Py_ssize_t ctx_in } } -static Py_ssize_t -mp_decode_map_size(DecoderState *self, char op, char *expected, TypeNode *ctx, Py_ssize_t ctx_ind) { +static PyObject * +mp_decode_type_dict(DecoderState *self, char op, TypeNodeMap *type, TypeNode *ctx, Py_ssize_t ctx_ind) { + Py_ssize_t size, i; + PyObject *res, *key = NULL, *val = NULL; + if ('\x80' <= op && op <= '\x8f') { - return op & 0x0f; + size = op & 0x0f; } else if (op == MP_MAP16) { - return mp_decode_size2(self); + size = mp_decode_size2(self); } else if (op == MP_MAP32) { - return mp_decode_size4(self); + size = mp_decode_size4(self); + } + else { + return mp_validation_error(op, "dict", ctx, ctx_ind); } - mp_validation_error(op, expected, ctx, ctx_ind); - return -1; -} - -static PyObject * -mp_decode_type_dict(DecoderState *self, char op, TypeNodeMap *type, TypeNode *ctx, Py_ssize_t ctx_ind) { - Py_ssize_t size, i; - PyObject *res, *key = NULL, *val = NULL; - - size = mp_decode_map_size(self, op, "dict", ctx, ctx_ind); if (size < 0) return NULL; res = PyDict_New(); @@ -4050,16 +4046,13 @@ mp_decode_cstr(DecoderState *self, char ** out, TypeNode *ctx, Py_ssize_t ctx_in } static PyObject * -mp_decode_type_struct(DecoderState *self, char op, TypeNodeObj *type, TypeNode *ctx, Py_ssize_t ctx_ind) { - Py_ssize_t i, size, key_size, field_index, nfields, ndefaults, pos = 0; +mp_decode_type_struct_map(DecoderState *self, Py_ssize_t size, TypeNodeObj *type, TypeNode *ctx, Py_ssize_t ctx_ind) { + Py_ssize_t i, key_size, field_index, nfields, ndefaults, pos = 0; char *key = NULL; PyObject *res, *val = NULL; StructMetaObject *st_type = (StructMetaObject *)(type->arg); int should_untrack; - size = mp_decode_map_size(self, op, "struct", ctx, ctx_ind); - if (size < 0) return NULL; - res = ((PyTypeObject *)(st_type))->tp_alloc((PyTypeObject *)st_type, 0); if (res == NULL) return NULL; @@ -4127,6 +4120,106 @@ mp_decode_type_struct(DecoderState *self, char op, TypeNodeObj *type, TypeNode * return NULL; } +static PyObject * +mp_decode_type_struct_array(DecoderState *self, Py_ssize_t size, TypeNodeObj *type, TypeNode *ctx, Py_ssize_t ctx_ind) { + Py_ssize_t i, nfields, ndefaults, npos; + PyObject *res, *val = NULL; + StructMetaObject *st_type = (StructMetaObject *)(type->arg); + int should_untrack; + + res = ((PyTypeObject *)(st_type))->tp_alloc((PyTypeObject *)st_type, 0); + if (res == NULL) return NULL; + + nfields = PyTuple_GET_SIZE(st_type->struct_fields); + ndefaults = PyTuple_GET_SIZE(st_type->struct_defaults); + npos = nfields - ndefaults; + should_untrack = PyObject_IS_GC(res); + + if (Py_EnterRecursiveCall(" while deserializing an object")) { + Py_DECREF(res); + return NULL; + } + for (i = 0; i < nfields; i++) { + if (size > 0) { + val = mp_decode_type(self, st_type->struct_types[i], (TypeNode *)type, i, false); + if (val == NULL) goto error; + size--; + } + else if (i < npos) { + PyErr_Format( + msgspec_get_global_state()->DecodingError, + "Error decoding `%s`: missing required field `%S`", + ((PyTypeObject *)st_type)->tp_name, + PyTuple_GET_ITEM(st_type->struct_fields, i) + ); + goto error; + } + else { + val = maybe_deepcopy_default( + PyTuple_GET_ITEM(st_type->struct_defaults, i - npos) + ); + if (val == NULL) + goto error; + } + Struct_set_index(res, i, val); + if (should_untrack) { + should_untrack = !OBJ_IS_GC(val); + } + } + /* Ignore all trailing fields */ + while (size > 0) { + if (mp_skip(self) < 0) + goto error; + size--; + } + Py_LeaveRecursiveCall(); + if (should_untrack) + PyObject_GC_UnTrack(res); + return res; +error: + Py_LeaveRecursiveCall(); + Py_DECREF(res); + return NULL; +} + +static PyObject * +mp_decode_type_struct(DecoderState *self, char op, TypeNodeObj *type, TypeNode *ctx, Py_ssize_t ctx_ind) { + Py_ssize_t size; + + if ('\x80' <= op && op <= '\x8f') { + size = op & 0x0f; + } + else if (op == MP_MAP16) { + size = mp_decode_size2(self); + } + else if (op == MP_MAP32) { + size = mp_decode_size4(self); + } + else if (((StructMetaObject *)type->arg)->asarray) { + if ('\x90' <= op && op <= '\x9f') { + size = op & 0x0f; + } + else if (op == MP_ARRAY16) { + size = mp_decode_size2(self); + } + else if (op == MP_ARRAY32) { + size = mp_decode_size4(self); + } + else { + return mp_validation_error(op, "struct", ctx, ctx_ind); + } + if (size < 0) return NULL; + return mp_decode_type_struct_array(self, size, type, ctx, ctx_ind); + } + else { + return mp_validation_error(op, "struct", ctx, ctx_ind); + } + + if (size < 0) return NULL; + return mp_decode_type_struct_map(self, size, type, ctx, ctx_ind); +} + + static Py_ssize_t mp_decode_array_size(DecoderState *self, char op, char *expected, TypeNode *ctx, Py_ssize_t ctx_ind) { if ('\x90' <= op && op <= '\x9f') { From e5eea16f7fb519b6e9aff6177cecc6ce86aa7976 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 11:38:04 -0500 Subject: [PATCH 5/6] Add tests for typed struct asarray decoding --- msgspec/core.c | 8 ++-- tests/test_msgspec.py | 94 +++++++++++++++++++++++++++++++++++++++---- tests/test_struct.py | 43 ++++++++++++-------- 3 files changed, 118 insertions(+), 27 deletions(-) diff --git a/msgspec/core.c b/msgspec/core.c index 045f8ccc..186ae832 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -3577,7 +3577,7 @@ static PyObject * mp_decode_type( static PyObject * mp_format_validation_error(const char *expected, const char *got, TypeNode *ctx, Py_ssize_t ctx_ind) { MsgspecState *st = msgspec_get_global_state(); - if (ctx->code == TYPE_STRUCT) { + if (ctx->code == TYPE_STRUCT && ctx_ind != -1) { StructMetaObject *st_type = (StructMetaObject *)(((TypeNodeObj *)ctx)->arg); PyObject *field = PyTuple_GET_ITEM(st_type->struct_fields, ctx_ind); PyObject *typstr = TypeNode_Repr(st_type->struct_types[ctx_ind]); @@ -4195,7 +4195,7 @@ mp_decode_type_struct(DecoderState *self, char op, TypeNodeObj *type, TypeNode * else if (op == MP_MAP32) { size = mp_decode_size4(self); } - else if (((StructMetaObject *)type->arg)->asarray) { + else if (((StructMetaObject *)type->arg)->asarray == OPT_TRUE) { if ('\x90' <= op && op <= '\x9f') { size = op & 0x0f; } @@ -4479,7 +4479,7 @@ Decoder_decode(Decoder *self, PyObject *const *args, Py_ssize_t nargs) self->state.input_buffer = buffer.buf; self->state.input_len = buffer.len; self->state.next_read_idx = 0; - res = mp_decode_type(&(self->state), self->state.type, self->state.type, 0, false); + res = mp_decode_type(&(self->state), self->state.type, self->state.type, -1, false); } if (buffer.buf != NULL) { @@ -4643,7 +4643,7 @@ msgspec_decode(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject state.input_len = buffer.len; state.next_read_idx = 0; if (state.type != NULL) { - res = mp_decode_type(&state, state.type, state.type, 0, false); + res = mp_decode_type(&state, state.type, state.type, -1, false); } else { res = mp_decode_any(&state, false); } diff --git a/tests/test_msgspec.py b/tests/test_msgspec.py index 64705917..5e6a3dd4 100644 --- a/tests/test_msgspec.py +++ b/tests/test_msgspec.py @@ -31,6 +31,13 @@ class Person(msgspec.Struct): prefect: bool = False +class PersonAA(msgspec.Struct, asarray=True): + first: str + last: str + age: int + prefect: bool = False + + class Node(msgspec.Struct): left: Optional[Node] = None right: Optional[Node] = None @@ -457,11 +464,12 @@ def test_decode_dec_hook_wrong_type(self): dec = msgspec.Decoder(type=Point, dec_hook=lambda t, o: o) buf = msgspec.encode((1, 2)) - with pytest.raises(msgspec.DecodingError) as rec: + with pytest.raises( + msgspec.DecodingError, + match="Error decoding `Point`: expected `Point`, got `list`", + ): dec.decode(buf) - assert "Error decoding `Point`: expected `Point`, got `list`" == str(rec.value) - def test_decode_dec_hook_wrong_type_in_struct(self): class Test(msgspec.Struct): point: Point @@ -864,12 +872,15 @@ def test_struct(self): with pytest.raises(msgspec.DecodingError, match="truncated"): dec.decode(a[:-2]) - with pytest.raises(msgspec.DecodingError, match="expected `struct`"): + with pytest.raises( + msgspec.DecodingError, + match="Error decoding `Person`: expected `struct`, got `int`", + ): dec.decode(enc.encode(1)) with pytest.raises( msgspec.DecodingError, - match=r"Error decoding `Person` field `first` \(`str`\): expected `str`, got `int`", + match="Error decoding `Person`: expected `str`, got `int`", ): dec.decode(enc.encode({1: "harry"})) @@ -937,8 +948,77 @@ def test_struct_defaults_missing_fields(self): assert res == Person("harry", "potter", 13) assert res.prefect is False - def test_struct_gc_maybe_untracked_on_decode(self): - class Test(msgspec.Struct): + def test_struct_asarray(self): + enc = msgspec.Encoder() + dec = msgspec.Decoder(PersonAA) + + x = PersonAA(first="harry", last="potter", age=13) + a = enc.encode(x) + assert enc.encode(("harry", "potter", 13, False)) == a + assert dec.decode(a) == x + + with pytest.raises(msgspec.DecodingError, match="truncated"): + dec.decode(a[:-2]) + + with pytest.raises( + msgspec.DecodingError, + match="Error decoding `PersonAA`: expected `struct`, got `int`", + ): + dec.decode(enc.encode(1)) + + # Wrong field type + bad = enc.encode(("harry", "potter", "thirteen")) + with pytest.raises(msgspec.DecodingError, match="expected `int`"): + dec.decode(bad) + + # Missing fields + bad = enc.encode(("harry", "potter")) + with pytest.raises(msgspec.DecodingError, match="missing required field `age`"): + dec.decode(bad) + + bad = enc.encode(()) + with pytest.raises( + msgspec.DecodingError, match="missing required field `first`" + ): + dec.decode(bad) + + # Extra fields ignored + dec2 = msgspec.Decoder(List[PersonAA]) + msg = enc.encode( + [ + ("harry", "potter", 13, False, 1, 2, 3, 4), + ("ron", "weasley", 13, False, 5, 6), + ] + ) + res = dec2.decode(msg) + assert res == [PersonAA("harry", "potter", 13), PersonAA("ron", "weasley", 13)] + + # Defaults applied + res = dec.decode(enc.encode(("harry", "potter", 13))) + assert res == PersonAA("harry", "potter", 13) + assert res.prefect is False + + def test_struct_only_asarray_structs_can_decode_from_array(self): + array_msg = msgspec.encode(("harry", "potter", 13)) + map_msg = msgspec.encode({"first": "harry", "last": "potter", "age": 13}) + sol = Person("harry", "potter", 13) + array_sol = PersonAA("harry", "potter", 13) + + dec = msgspec.Decoder(Person) + array_dec = msgspec.Decoder(PersonAA) + + assert array_dec.decode(map_msg) == array_sol + assert array_dec.decode(array_msg) == array_sol + assert dec.decode(map_msg) == sol + with pytest.raises( + msgspec.DecodingError, + match="Error decoding `Person`: expected `struct`, got `list`", + ): + dec.decode(array_msg) + + @pytest.mark.parametrize("asarray", [False, True]) + def test_struct_gc_maybe_untracked_on_decode(self, asarray): + class Test(msgspec.Struct, asarray=asarray): x: Any y: Any z: Tuple = () diff --git a/tests/test_struct.py b/tests/test_struct.py index 66486b4c..88c0ad59 100644 --- a/tests/test_struct.py +++ b/tests/test_struct.py @@ -647,46 +647,57 @@ def test_struct_handles_missing_attributes(): pickle.dumps(t) -def test_struct_option_precedence(): +@pytest.mark.parametrize("option", ["immutable", "asarray"]) +def test_struct_option_precedence(option): + def get(cls): + return getattr(cls, option) + class Default(Struct): pass - assert not Default.immutable + assert not get(Default) - class Immutable(Struct, immutable=True): + class Enabled(Struct, **{option: True}): pass - assert Immutable.immutable + assert get(Enabled) - class NotImmutable(Struct, immutable=False): + class Disabled(Struct, **{option: False}): pass - assert not NotImmutable.immutable + assert not get(Disabled) - class T(Immutable): + class T(Enabled): pass - assert T.immutable + assert get(T) - class T(Immutable, immutable=False): + class T(Enabled, **{option: False}): pass - assert not T.immutable + assert not get(T) - class T(Immutable, Default): + class T(Enabled, Default): pass - assert T.immutable + assert get(T) - class T(Default, Immutable): + class T(Default, Enabled): pass - assert T.immutable + assert get(T) - class T(Default, NotImmutable, Immutable): + class T(Default, Disabled, Enabled): pass - assert not T.immutable + assert not get(T) + + +def test_invalid_option_raises(): + with pytest.raises(TypeError): + + class Foo(Struct, invalid=True): + pass class ImmutablePoint(Struct, immutable=True): From a82f6bb983842cb46f0f11bedaa653f502b1b584 Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Sat, 3 Jul 2021 12:08:24 -0500 Subject: [PATCH 6/6] Add docs --- docs/source/index.rst | 64 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 54 insertions(+), 10 deletions(-) diff --git a/docs/source/index.rst b/docs/source/index.rst index 01ace398..fafba10d 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -202,19 +202,22 @@ acceptible: - `list` / `typing.List` - `dict` / `typing.Dict` - `set` / `typing.Set` +- `datetime.datetime` - `typing.Any` - `typing.Optional` - `msgspec.Ext` - `enum.Enum` derived types - `enum.IntEnum` derived types - `msgspec.Struct` derived types +- Custom types (provided a valid ``ext_hook`` or ``dec_hook`` callback on the Decoder) + Structs ~~~~~~~ ``msgspec`` can serialize many builtin types, but unlike protocols like -`pickle`_/`quickle`_, it can't serialize arbitrary user classes. Two -user-defined types are supported: +`pickle`_/`quickle`_, it can't serialize arbitrary user classes by default. Two +user-defined types are supported though: - `Struct` - `enum.Enum` @@ -254,11 +257,31 @@ annotations: >>> ron == harry False -It is forbidden to override ``__init__``/``__new__`` in a struct definition, -but other methods can be overridden or added as needed. The struct fields are -available via the ``__struct_fields__`` attribute (a tuple of the fields in -argument order ) if you need them. Here we add a method for converting a struct -to a dict. +If needed, a ``__hash__`` method can also be generated by specifying +``immutable=True`` when defining the struct. Note that this disables modifying +field values after initialization. + +.. code-block:: python + + >>> class Point(msgspec.Struct, immutable=True): + ... """This struct is immutable & hashable""" + ... x: float + ... y: float + ... + >>> p = Point(1.0, 2.0) + >>> {p: 1} # immutable structs are hashable, and can be keys in dicts + {Point(1.0, 2.0): 1} + >>> p.x = 2.0 # immutable structs cannot be modified after creation + Traceback (most recent call last): + ... + TypeError: immutable type: 'Point' + +Note that it is forbidden to override ``__init__``/``__new__`` in a struct +definition, but other methods can be overridden or added as needed. + +The struct fields are available via the ``__struct_fields__`` attribute (a +tuple of the fields in argument order ) if you need them. Here we add a method +for converting a struct to a dict. .. code-block:: python @@ -295,6 +318,25 @@ deserialization, it also can improve performance. Depending on the schema, deserializing a message into a `Struct` can be *roughly twice as fast* as deserializing it into a `dict`. +If you need higher performance (at the cost of more inscrutable message +encoding), you can set ``asarray=True`` on a struct definition. Structs with +this option enabled are encoded/decoded as MessagePack ``array`` types (rather +than ``map`` types), removing the field names from the serialized message. This +can provide another ~2x speedup for decoding (and ~1.5x speedup for encoding). + +.. code-block:: python + + >>> class ArrayBasedStruct(msgspec.Struct, asarray=True): + ... """This struct is serialized as a MessagePack array type + ... (instead of a map type). This means no field names are sent + ... as part of the message, speeding up encoding/decoding.""" + ... my_first_field: str + ... my_second_field: int + ... + >>> x = ArrayBasedStruct("some string", 2) + >>> msgspec.encode(x) + b'\x92\xabsome string\x02' + .. _extensions: Extensions @@ -436,9 +478,11 @@ mismatched versions. For schema evolution to work smoothly, you need to follow a few guidelines: 1. Any new fields on a `Struct` must specify default values. -2. Don't change the type annotations for existing messages or fields -3. Don't change the type codes or implementations for any defined - :ref:`Extensions` +2. Structs with ``asarray=True`` must not reorder fields, and any new fields + must be appended to the end (and have defaults). +3. Don't change the type annotations for existing messages or fields. +4. Don't change the type codes or implementations for any defined + :ref:`Extensions`. For example, suppose we wanted to add a new ``email`` field to our ``Person`` struct. To do so, we add it at the end of the definition, with a default value.