gh-129559: Add bytearray.resize() (GH-129560)

Add bytearray.resize() which wraps PyByteArray_Resize.

Make negative size passed to resize exception/error rather than crash in optimized builds.
This commit is contained in:
Cody Maloney 2025-02-05 11:33:17 -08:00 committed by GitHub
parent 7d9a22f509
commit 5fb019fc29
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 158 additions and 9 deletions

View file

@ -74,6 +74,11 @@ Direct API functions
.. c:function:: int PyByteArray_Resize(PyObject *bytearray, Py_ssize_t len)
Resize the internal buffer of *bytearray* to *len*.
Failure is a ``-1`` return with an exception set.
.. versionchanged:: next
A negative *len* will now result in an exception being set and -1 returned.
Macros
^^^^^^

View file

@ -2841,6 +2841,38 @@ objects.
optional *sep* and *bytes_per_sep* parameters to insert separators
between bytes in the hex output.
.. method:: resize(size)
Resize the :class:`bytearray` to contain *size* bytes. *size* must be
greater than or equal to 0.
If the :class:`bytearray` needs to shrink, bytes beyond *size* are truncated.
If the :class:`bytearray` needs to grow, all new bytes, those beyond *size*,
will be set to null bytes.
This is equivalent to:
>>> def resize(ba, size):
... if len(ba) > size:
... del ba[size:]
... else:
... ba += b'\0' * (size - len(ba))
Examples:
>>> shrink = bytearray(b'abc')
>>> shrink.resize(1)
>>> (shrink, len(shrink))
(bytearray(b'a'), 1)
>>> grow = bytearray(b'abc')
>>> grow.resize(5)
>>> (grow, len(grow))
(bytearray(b'abc\x00\x00'), 5)
.. versionadded:: next
Since bytearray objects are sequences of integers (akin to a list), for a
bytearray object *b*, ``b[0]`` will be an integer, while ``b[0:1]`` will be
a bytearray object of length 1. (This contrasts with text strings, where

View file

@ -1359,6 +1359,44 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
b = by("Hello, world")
self.assertEqual(re.findall(br"\w+", b), [by("Hello"), by("world")])
def test_resize(self):
ba = bytearray(b'abcdef')
self.assertIsNone(ba.resize(3))
self.assertEqual(ba, bytearray(b'abc'))
self.assertIsNone(ba.resize(10))
self.assertEqual(len(ba), 10)
# Bytes beyond set values must be cleared.
self.assertEqual(ba, bytearray(b'abc\0\0\0\0\0\0\0'))
ba[3:10] = b'defghij'
self.assertEqual(ba, bytearray(b'abcdefghij'))
self.assertIsNone(ba.resize(2 ** 20))
self.assertEqual(len(ba), 2**20)
self.assertEqual(ba, bytearray(b'abcdefghij' + b'\0' * (2 ** 20 - 10)))
self.assertIsNone(ba.resize(0))
self.assertEqual(ba, bytearray())
self.assertIsNone(ba.resize(10))
self.assertEqual(ba, bytearray(b'\0' * 10))
# Subclass
ba = ByteArraySubclass(b'abcdef')
self.assertIsNone(ba.resize(3))
self.assertEqual(ba, bytearray(b'abc'))
# Check arguments
self.assertRaises(TypeError, bytearray().resize)
self.assertRaises(TypeError, bytearray().resize, (10, 10))
self.assertRaises(ValueError, bytearray().resize, -1)
self.assertRaises(ValueError, bytearray().resize, -200)
self.assertRaises(MemoryError, bytearray().resize, sys.maxsize)
self.assertRaises(MemoryError, bytearray(1000).resize, sys.maxsize)
def test_setitem(self):
def setitem_as_mapping(b, i, val):
b[i] = val
@ -1715,17 +1753,18 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
# if it wouldn't reallocate the underlying buffer.
# Furthermore, no destructive changes to the buffer may be applied
# before raising the error.
b = bytearray(range(10))
b = bytearray(10)
v = memoryview(b)
def resize(n):
def manual_resize(n):
b[1:-1] = range(n + 1, 2*n - 1)
resize(10)
b.resize(10)
orig = b[:]
self.assertRaises(BufferError, resize, 11)
self.assertRaises(BufferError, b.resize, 11)
self.assertRaises(BufferError, manual_resize, 11)
self.assertEqual(b, orig)
self.assertRaises(BufferError, resize, 9)
self.assertRaises(BufferError, b.resize, 9)
self.assertEqual(b, orig)
self.assertRaises(BufferError, resize, 0)
self.assertRaises(BufferError, b.resize, 0)
self.assertEqual(b, orig)
# Other operations implying resize
self.assertRaises(BufferError, b.pop, 0)

View file

@ -151,10 +151,11 @@ class CAPITest(unittest.TestCase):
self.assertEqual(resize(ba, 3), 0)
self.assertEqual(ba, bytearray(b'abc'))
self.assertRaises(ValueError, resize, bytearray(), -1)
self.assertRaises(ValueError, resize, bytearray(), -200)
self.assertRaises(MemoryError, resize, bytearray(), PY_SSIZE_T_MAX)
self.assertRaises(MemoryError, resize, bytearray(1000), PY_SSIZE_T_MAX)
# CRASHES resize(bytearray(b'abc'), -1)
# CRASHES resize(b'abc', 0)
# CRASHES resize(object(), 0)
# CRASHES resize(NULL, 0)

View file

@ -0,0 +1,2 @@
Add :meth:`bytearray.resize` method so :class:`bytearray` can be efficiently
resized in place.

View file

@ -184,7 +184,12 @@ PyByteArray_Resize(PyObject *self, Py_ssize_t requested_size)
assert(self != NULL);
assert(PyByteArray_Check(self));
assert(logical_offset <= alloc);
assert(requested_size >= 0);
if (requested_size < 0) {
PyErr_Format(PyExc_ValueError,
"Can only resize to positive sizes, got %zd", requested_size);
return -1;
}
if (requested_size == Py_SIZE(self)) {
return 0;
@ -1388,6 +1393,31 @@ bytearray_removesuffix_impl(PyByteArrayObject *self, Py_buffer *suffix)
}
/*[clinic input]
bytearray.resize
size: Py_ssize_t
New size to resize to..
/
Resize the internal buffer of bytearray to len.
[clinic start generated code]*/
static PyObject *
bytearray_resize_impl(PyByteArrayObject *self, Py_ssize_t size)
/*[clinic end generated code: output=f73524922990b2d9 input=75fd4d17c4aa47d3]*/
{
Py_ssize_t start_size = PyByteArray_GET_SIZE(self);
int result = PyByteArray_Resize((PyObject *)self, size);
if (result < 0) {
return NULL;
}
// Set new bytes to null bytes
if (size > start_size) {
memset(PyByteArray_AS_STRING(self) + start_size, 0, size - start_size);
}
Py_RETURN_NONE;
}
/*[clinic input]
bytearray.translate
@ -2361,6 +2391,7 @@ static PyMethodDef bytearray_methods[] = {
BYTEARRAY_REPLACE_METHODDEF
BYTEARRAY_REMOVEPREFIX_METHODDEF
BYTEARRAY_REMOVESUFFIX_METHODDEF
BYTEARRAY_RESIZE_METHODDEF
BYTEARRAY_REVERSE_METHODDEF
BYTEARRAY_RFIND_METHODDEF
BYTEARRAY_RINDEX_METHODDEF

View file

@ -565,6 +565,45 @@ exit:
return return_value;
}
PyDoc_STRVAR(bytearray_resize__doc__,
"resize($self, size, /)\n"
"--\n"
"\n"
"Resize the internal buffer of bytearray to len.\n"
"\n"
" size\n"
" New size to resize to..");
#define BYTEARRAY_RESIZE_METHODDEF \
{"resize", (PyCFunction)bytearray_resize, METH_O, bytearray_resize__doc__},
static PyObject *
bytearray_resize_impl(PyByteArrayObject *self, Py_ssize_t size);
static PyObject *
bytearray_resize(PyObject *self, PyObject *arg)
{
PyObject *return_value = NULL;
Py_ssize_t size;
{
Py_ssize_t ival = -1;
PyObject *iobj = _PyNumber_Index(arg);
if (iobj != NULL) {
ival = PyLong_AsSsize_t(iobj);
Py_DECREF(iobj);
}
if (ival == -1 && PyErr_Occurred()) {
goto exit;
}
size = ival;
}
return_value = bytearray_resize_impl((PyByteArrayObject *)self, size);
exit:
return return_value;
}
PyDoc_STRVAR(bytearray_translate__doc__,
"translate($self, table, /, delete=b\'\')\n"
"--\n"
@ -1623,4 +1662,4 @@ bytearray_sizeof(PyObject *self, PyObject *Py_UNUSED(ignored))
{
return bytearray_sizeof_impl((PyByteArrayObject *)self);
}
/*[clinic end generated code: output=bc8bec8514102bf3 input=a9049054013a1b77]*/
/*[clinic end generated code: output=41bb67a8a181e733 input=a9049054013a1b77]*/