msgspec · jcrist · Jun 28, 2021 · Jun 28, 2021 · Jun 28, 2021
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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):
-Original file line number
+Diff line change
@@ Expand Up / @@ -7,7 +7,7 @@ Encoder @@
     -------
     .. autoclass:: Encoder
-        :members: encode
+        :members: encode, encode_into
     Decoder
@@ Expand Down @@