Skip to content

Commit

Permalink
Change behavior of encode_into with offset > len(buf)
Browse files Browse the repository at this point in the history
Previously we'd truncate `offset = len(buf)` if `offset > len(buf)` when
calling `encode_into(msg, buf, offset)`. We now expand the size of `buf`
to `offset`, providing a more consistent and expected behavior.
  • Loading branch information
jcrist committed Oct 20, 2024
1 parent 150f64a commit 7745144
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 8 deletions.
8 changes: 6 additions & 2 deletions msgspec/_core.c
Original file line number Diff line number Diff line change
Expand Up @@ -9401,8 +9401,10 @@ encoder_encode_into_common(
PyErr_SetString(PyExc_ValueError, "offset must be >= -1");
return NULL;
}
if (offset > buf_size) {
offset = buf_size;

if (offset < buf_size) {
buf_size = Py_MAX(8, 1.5 * offset);
if (PyByteArray_Resize(buf, buf_size) < 0) return NULL;
}
}

Expand All @@ -9419,9 +9421,11 @@ encoder_encode_into_common(
.max_output_len = buf_size,
.resize_buffer = ms_resize_bytearray
};

if (encode(&state, obj) < 0) {
return NULL;
}

FAST_BYTEARRAY_SHRINK(buf, state.output_len);
Py_RETURN_NONE;
}
Expand Down
7 changes: 4 additions & 3 deletions tests/test_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -298,10 +298,11 @@ def test_encode_into_offset(self):
enc.encode_into(msg, buf, 2)
assert buf == b"01" + encoded

# Offset out of bounds appends to end
# Offset out of bounds extends
buf = bytearray(b"01234")
enc.encode_into(msg, buf, 1000)
assert buf == b"01234" + encoded
enc.encode_into(msg, buf, 10)
assert buf[:5] == b"01234"
assert buf[10:] == encoded

# Offset -1 means append at end
buf = bytearray(b"01234")
Expand Down
7 changes: 4 additions & 3 deletions tests/test_msgpack.py
Original file line number Diff line number Diff line change
Expand Up @@ -432,10 +432,11 @@ def test_encode_into_offset(self):
enc.encode_into(msg, buf, 2)
assert buf == b"01" + encoded

# Offset out of bounds appends to end
# Offset out of bounds extends
buf = bytearray(b"01234")
enc.encode_into(msg, buf, 1000)
assert buf == b"01234" + encoded
enc.encode_into(msg, buf, 10)
assert buf[:5] == b"01234"
assert buf[10:] == encoded

# Offset -1 means append at end
buf = bytearray(b"01234")
Expand Down

0 comments on commit 7745144

Please sign in to comment.