diff --git a/docs/source/api.rst b/docs/source/api.rst index 7f0c7772..ac0d1bac 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -7,7 +7,7 @@ Encoder ------- .. autoclass:: Encoder - :members: encode + :members: encode, encode_into Decoder diff --git a/msgspec/core.c b/msgspec/core.c index 0ccc28e4..902f5921 100644 --- a/msgspec/core.c +++ b/msgspec/core.c @@ -9,9 +9,11 @@ #if PY_VERSION_HEX < 0x03090000 #define IS_TRACKED _PyObject_GC_IS_TRACKED #define CALL_ONE_ARG(fn, arg) PyObject_CallFunctionObjArgs((fn), (arg), NULL) +#define SET_SIZE(obj, size) (((PyVarObject *)obj)->ob_size = size) #else #define IS_TRACKED PyObject_GC_IsTracked #define CALL_ONE_ARG(fn, arg) PyObject_CallOneArg((fn), (arg)) +#define SET_SIZE(obj, size) Py_SET_SIZE(obj, size) #endif /* Is this object something that is/could be GC tracked? True if * - the value supports GC @@ -21,8 +23,6 @@ (PyType_IS_GC(Py_TYPE(x)) && \ (!PyTuple_CheckExact(x) || IS_TRACKED(x))) - - /************************************************************************* * Endian handling macros * *************************************************************************/ @@ -1712,7 +1712,7 @@ typedef struct EncoderState { PyObject *enc_hook; /* `enc_hook` callback */ Py_ssize_t write_buffer_size; /* Configured internal buffer size */ - PyObject *output_buffer; /* Bytearray storing the output */ + PyObject *output_buffer; /* bytes or bytearray storing the output */ Py_ssize_t output_len; /* Length of output_buffer */ Py_ssize_t max_output_len; /* Allocation size of output_buffer */ } EncoderState; @@ -1845,19 +1845,25 @@ mp_write(EncoderState *self, const char *s, Py_ssize_t n) { Py_ssize_t required; char *buffer; + bool is_bytes = PyBytes_CheckExact(self->output_buffer); required = self->output_len + n; if (required > self->max_output_len) { /* Make space in buffer */ - if (self->output_len >= PY_SSIZE_T_MAX / 2 - n) { - PyErr_NoMemory(); - return -1; - } - self->max_output_len = (self->output_len + n) / 2 * 3; - if (_PyBytes_Resize(&self->output_buffer, self->max_output_len) < 0) - return -1; + int status; + self->max_output_len = Py_MAX(8, (self->output_len + n) / 2 * 3); + status = ( + is_bytes ? _PyBytes_Resize(&self->output_buffer, self->max_output_len) + : PyByteArray_Resize(self->output_buffer, self->max_output_len) + ); + if (status < 0) return -1; + } + if (is_bytes) { + buffer = PyBytes_AS_STRING(self->output_buffer); + } + else { + buffer = PyByteArray_AS_STRING(self->output_buffer); } - buffer = PyBytes_AS_STRING(self->output_buffer); memcpy(buffer + self->output_len, s, n); self->output_len += n; return 0; @@ -2415,6 +2421,92 @@ mp_encode(EncoderState *self, PyObject *obj) } } +PyDoc_STRVAR(Encoder_encode_into__doc__, +"encode_into(self, obj, buffer, offset=0, /)\n" +"--\n" +"\n" +"Serialize an object into an existing bytearray buffer.\n" +"\n" +"Upon success, the buffer will be truncated to the end of the serialized\n" +"message. Note that the underlying memory buffer *won't* be truncated,\n" +"allowing for efficiently appending additional bytes later.\n" +"\n" +"Parameters\n" +"----------\n" +"obj : Any\n" +" The object to serialize.\n" +"buffer : bytearray\n" +" The buffer to serialize into.\n" +"offset : int, optional\n" +" The offset into the buffer to start writing at. Defaults to 0. Set to -1\n" +" to start writing at the end of the buffer.\n" +"\n" +"Returns\n" +"-------\n" +"None" +); +static PyObject* +Encoder_encode_into(Encoder *self, PyObject *const *args, Py_ssize_t nargs) +{ + int status; + PyObject *obj, *old_buf, *buf; + Py_ssize_t buf_size, offset = 0; + + if (!check_positional_nargs(nargs, 2, 3)) { + return NULL; + } + obj = args[0]; + buf = args[1]; + if (!PyByteArray_CheckExact(buf)) { + PyErr_SetString(PyExc_TypeError, "buffer must be a `bytearray`"); + return NULL; + } + buf_size = PyByteArray_GET_SIZE(buf); + + if (nargs == 3) { + offset = PyLong_AsSsize_t(args[2]); + if (offset == -1) { + if (PyErr_Occurred()) return NULL; + offset = buf_size; + } + if (offset < 0) { + PyErr_SetString(PyExc_ValueError, "offset must be >= -1"); + return NULL; + } + if (offset > buf_size) { + offset = buf_size; + } + } + + /* Setup buffer */ + old_buf = self->state.output_buffer; + self->state.output_buffer = buf; + self->state.output_len = offset; + self->state.max_output_len = buf_size; + + status = mp_encode(&(self->state), obj); + + if (status == 0) { + /* Set the length of the bytearray *without* actually resizing the + * backing memory buffer. This is useful for propagating size info + * downstream without doing any additional memory operations. Most + * users of this method will either immediately write more onto the end + * of the buffer (in which case they will need to resize the buffer + * back up anyway), or they will write the buffer to a socket/file/... + * and release the memory. Either way, hitting realloc here seems + * unnecessary. + * + * This is copied from within the fastpath of `PyByteArray_Resize`*/ + SET_SIZE(self->state.output_buffer, self->state.output_len); + PyByteArray_AS_STRING(self->state.output_buffer)[self->state.output_len] = '\0'; + } + + /* Reset buffer */ + self->state.output_buffer = old_buf; + + Py_RETURN_NONE; +} + PyDoc_STRVAR(Encoder_encode__doc__, "encode(self, obj)\n" "--\n" @@ -2495,6 +2587,10 @@ static struct PyMethodDef Encoder_methods[] = { "encode", (PyCFunction) Encoder_encode, METH_FASTCALL, Encoder_encode__doc__, }, + { + "encode_into", (PyCFunction) Encoder_encode_into, METH_FASTCALL, + Encoder_encode_into__doc__, + }, { "__sizeof__", (PyCFunction) Encoder_sizeof, METH_NOARGS, PyDoc_STR("Size in bytes") diff --git a/msgspec/core.pyi b/msgspec/core.pyi index 5f62e3b7..7f32e831 100644 --- a/msgspec/core.pyi +++ b/msgspec/core.pyi @@ -37,6 +37,9 @@ class Encoder: write_buffer_size: int = ..., ): ... def encode(self, obj: Any) -> bytes: ... + def encode_into( + self, obj: Any, buffer: bytearray, offset: Optional[int] = 0 + ) -> None: ... @overload def decode(buf: bytes, ext_hook: ext_hook_sig = None) -> Any: ... diff --git a/tests/mypy_examples.py b/tests/mypy_examples.py index 7f17074c..2f836154 100644 --- a/tests/mypy_examples.py +++ b/tests/mypy_examples.py @@ -11,6 +11,13 @@ def check_Encoder_encode() -> None: reveal_type(b) # assert "bytes" in typ +def check_Encoder_encode_into() -> None: + enc = msgspec.Encoder() + buf = bytearray(48) + enc.encode_into([1, 2, 3], buf) + enc.encode_into([1, 2, 3], buf, 2) + + def check_encode() -> None: b = msgspec.encode([1, 2, 3]) diff --git a/tests/test_msgspec.py b/tests/test_msgspec.py index b1e2990e..45c18409 100644 --- a/tests/test_msgspec.py +++ b/tests/test_msgspec.py @@ -306,6 +306,70 @@ def test_encode_enc_hook_recursion_error(self): with pytest.raises(RecursionError): enc.encode(object()) + def test_encode_into_bad_arguments(self): + enc = msgspec.Encoder() + + with pytest.raises(TypeError, match="bytearray"): + enc.encode_into(1, b"test") + + with pytest.raises(TypeError): + enc.encode_into(1, bytearray(), "bad") + + with pytest.raises(ValueError, match="offset"): + enc.encode_into(1, bytearray(), -2) + + @pytest.mark.parametrize("buf_size", [0, 1, 16, 55, 60]) + def test_encode_into(self, buf_size): + enc = msgspec.Encoder() + + msg = {"key": "x" * 48} + encoded = msgspec.encode(msg) + + buf = bytearray(buf_size) + out = enc.encode_into(msg, buf) + assert out is None + assert buf == encoded + + def test_encode_into_offset(self): + enc = msgspec.Encoder() + msg = {"key": "value"} + encoded = enc.encode(msg) + + # Offset 0 is default + buf = bytearray() + enc.encode_into(msg, buf, 0) + assert buf == encoded + + # Offset in bounds uses the provided offset + buf = bytearray(b"01234") + enc.encode_into(msg, buf, 2) + assert buf == b"01" + encoded + + # Offset out of bounds appends to end + buf = bytearray(b"01234") + enc.encode_into(msg, buf, 1000) + assert buf == b"01234" + encoded + + # Offset -1 means append at end + buf = bytearray(b"01234") + enc.encode_into(msg, buf, -1) + assert buf == b"01234" + encoded + + def tesT_encode_into_handles_errors_properly(self): + enc = msgspec.Encoder() + out1 = enc.encode([1, 2, 3]) + + msg = [1, 2, object()] + buf = bytearray() + with pytest.raises(TypeError): + enc.encode_into(msg, buf) + + assert buf # buffer isn't reset upon error + + # Encoder still works + out2 = enc.encode([1, 2, 3]) + assert out1 == out2 + class TestDecoderMisc: def test_decoder_type_attribute(self):