add memoryview for decode

jcrist · Jan 5, 2024 · f3c13ec · f3c13ec
1 parent 4d4a02e
commit f3c13ec
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/msgspec/_core.c b/msgspec/_core.c
@@ -2525,6 +2525,7 @@ static PyTypeObject Field_Type = {
 #define MS_TYPE_TYPEDDICT           (1ull << 32)
 #define MS_TYPE_DATACLASS           (1ull << 33)
 #define MS_TYPE_NAMEDTUPLE          (1ull << 34)
+#define MS_TYPE_MEMORYVIEW          (1ull << 35)
 /* Constraints */
 #define MS_CONSTR_INT_MIN           (1ull << 42)
 #define MS_CONSTR_INT_MAX           (1ull << 43)
@@ -4623,6 +4624,10 @@ typenode_collect_type(TypeNodeCollectState *state, PyObject *obj) {
         state->types |= MS_TYPE_BYTEARRAY;
         kind = CK_BYTES;
     }
+    else if (t == (PyObject *)(&PyMemoryView_Type)) {
+        state->types |= MS_TYPE_MEMORYVIEW;
+        kind = CK_BYTES;
+    }
     else if (t == (PyObject *)(PyDateTimeAPI->DateTimeType)) {
         state->types |= MS_TYPE_DATETIME;
         kind = CK_TIME;
@@ -14464,6 +14469,9 @@ mpack_decode_bin(
     else if (type->types & MS_TYPE_UUID) {
         return ms_decode_uuid_from_bytes(s, size, path);
     }
+    else if (type->types & MS_TYPE_MEMORYVIEW) {
+        return PyMemoryView_FromMemory(s, size, PyBUF_READ);
+    }
 
     return ms_validation_error("bytes", type, path);
 }

diff --git a/tests/test_msgpack.py b/tests/test_msgpack.py
@@ -466,6 +466,13 @@ def test_encode_memoryview(self, size):
         res = msgspec.msgpack.decode(msg)
         assert buf == res
 
+    @pytest.mark.parametrize("size", SIZES)
+    def test_decode_memoryview(self, size):
+        buf = bytearray(size)
+        msg = msgspec.msgpack.encode(memoryview(buf))
+        res = msgspec.msgpack.decode(msg, type=memoryview)
+        assert buf == res.tobytes()
+
     @pytest.mark.parametrize(
         "dt, dt_str",
         [