gh-139871: Add bytearray.take_bytes([n]) to efficiently extract bytes

Doc/library/stdtypes.rst

@@ -3173,6 +3173,30 @@ objects.
 
  .. versionadded:: 3.14
 
+ .. method:: take_bytes(n=None, /)
+
+ Remove the first *n* bytes from the bytearray and return them as an immutable
+ :class:`bytes`.
+ By default (if *n* is ``None``), return all bytes and clear the bytearray.
+
+ If *n* is negative, index from the end and take the first :func:`len`
+ plus *n* bytes. If *n* is out of bounds, raise :exc:`IndexError`.
+
+ Taking less than the full length will leave remaining bytes in the
+ :class:`bytearray`, which requires a copy. If the remaining bytes should be
+ discarded, use :func:`~bytearray.resize` or :keyword:`del` to truncate
+ then :func:`~bytearray.take_bytes` without a size.
+
+ .. impl-detail::
+
+ Taking all bytes is a zero-copy operation.
+
+ .. versionadded:: next
+
+ See the :ref:`What's New <whatsnew315-bytearray-take-bytes>` entry for
+ common code patterns which can be optimized with
+ :func:`bytearray.take_bytes`.
+
 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

Doc/whatsnew/3.15.rst

@@ -307,6 +307,86 @@ Other language changes
  not only integers or floats, although this does not improve precision.
  (Contributed by Serhiy Storchaka in :gh:`67795`.)
 
+.. _whatsnew315-bytearray-take-bytes:
+
+* Added :meth:`bytearray.take_bytes(n=None, /) <bytearray.take_bytes>` to take
+ bytes out of a :class:`bytearray` without copying. This enables optimizing code
+ which must return :class:`bytes` after working with a mutable buffer of bytes
+ such as data buffering, network protocol parsing, encoding, decoding,
+ and compression. Common code patterns which can be optimized with
+ :func:`~bytearray.take_bytes` are listed below.
+
+ (Contributed by Cody Maloney in :gh:`139871`.)
+
+ .. list-table:: Suggested Optimizing Refactors
+ :header-rows: 1
+
+ * - Description
+ - Old
+ - New
+
+ * - Return :class:`bytes` after working with :class:`bytearray`
+ - .. code:: python
+
+ def read() -> bytes:
+ buffer = bytearray(1024)
+ ...
+ return bytes(buffer)
+
+ - .. code:: python
+
+ def read() -> bytes:
+ buffer = bytearray(1024)
+ ...
+ return buffer.take_bytes()
+
+ * - Empty a buffer getting the bytes
+ - .. code:: python
+
+ buffer = bytearray(1024)
+ ...
+ data = bytes(buffer)
+ buffer.clear()
+
+ - .. code:: python
+
+ buffer = bytearray(1024)
+ ...
+ data = buffer.take_bytes()
+
+ * - Split a buffer at a specific separator
+ - .. code:: python
+
+ buffer = bytearray(b'abc\ndef')
+ n = buffer.find(b'\n')
+ data = bytes(buffer[:n + 1])
+ del buffer[:n + 1]
+ assert data == b'abc'
+ assert buffer == bytearray(b'def')
+
+ - .. code:: python
+
+ buffer = bytearray(b'abc\ndef')
+ n = buffer.find(b'\n')
+ data = buffer.take_bytes(n + 1)
+
+ * - Split a buffer at a specific separator; discard after the separator
+ - .. code:: python
+
+ buffer = bytearray(b'abc\ndef')
+ n = buffer.find(b'\n')
+ data = bytes(buffer[:n])
+ buffer.clear()
+ assert data == b'abc'
+ assert len(buffer) == 0
+
+ - .. code:: python
+
+ buffer = bytearray(b'abc\ndef')
+ n = buffer.find(b'\n')
+ buffer.resize(n)
+ data = buffer.take_bytes()
+
 * Many functions related to compiling or parsing Python code, such as
  :func:`compile`, :func:`ast.parse`, :func:`symtable.symtable`,
  and :func:`importlib.abc.InspectLoader.source_to_code`, now allow to pass

Include/cpython/bytearrayobject.h

@@ -5,25 +5,25 @@
 /* Object layout */
 typedef struct {
  PyObject_VAR_HEAD
- Py_ssize_t ob_alloc; /* How many bytes allocated in ob_bytes */
+ /* How many bytes allocated in ob_bytes
+
+ In the current implementation this is equivalent to Py_SIZE(ob_bytes_object).
+ The value is always loaded and stored atomically for thread safety.
+ There are API compatibilty concerns with removing so keeping for now. */
+ Py_ssize_t ob_alloc;
  char *ob_bytes; /* Physical backing buffer */
  char *ob_start; /* Logical start inside ob_bytes */
  Py_ssize_t ob_exports; /* How many buffer exports */
+ PyObject *ob_bytes_object; /* PyBytes for zero-copy bytes conversion */
 } PyByteArrayObject;
 
-PyAPI_DATA(char) _PyByteArray_empty_string[];
-
 /* Macros and static inline functions, trading safety for speed */
 #define _PyByteArray_CAST(op) \
  (assert(PyByteArray_Check(op)), _Py_CAST(PyByteArrayObject*, op))
 
 static inline char* PyByteArray_AS_STRING(PyObject *op)
 {
- PyByteArrayObject *self = _PyByteArray_CAST(op);
- if (Py_SIZE(self)) {
- return self->ob_start;
- }
- return _PyByteArray_empty_string;
+ return _PyByteArray_CAST(op)->ob_start;
 }
 #define PyByteArray_AS_STRING(self) PyByteArray_AS_STRING(_PyObject_CAST(self))
 

Include/internal/pycore_bytesobject.h

@@ -60,6 +60,14 @@ PyAPI_FUNC(void)
 _PyBytes_Repeat(char* dest, Py_ssize_t len_dest,
  const char* src, Py_ssize_t len_src);
 
+/* _PyBytesObject_SIZE gives the basic size of a bytes object; any memory allocation
+ for a bytes object of length n should request PyBytesObject_SIZE + n bytes.
+
+ Using _PyBytesObject_SIZE instead of sizeof(PyBytesObject) saves
+ 3 or 7 bytes per bytes object allocation on a typical system.
+*/
+#define _PyBytesObject_SIZE (offsetof(PyBytesObject, ob_sval) + 1)
+
 /* --- PyBytesWriter ------------------------------------------------------ */
 
 struct PyBytesWriter {

Lib/test/test_bytes.py

@@ -1397,6 +1397,16 @@ def test_clear(self):
  b.append(ord('p'))
  self.assertEqual(b, b'p')
 
+ # Cleared object should be empty.
+ b = bytearray(b'abc')
+ b.clear()
+ self.assertEqual(b.__alloc__(), 0)
+ base_size = sys.getsizeof(bytearray())
+ self.assertEqual(sys.getsizeof(b), base_size)
+ c = b.copy()
+ self.assertEqual(c.__alloc__(), 0)
+ self.assertEqual(sys.getsizeof(c), base_size)
+
  def test_copy(self):
  b = bytearray(b'abc')
  bb = b.copy()
@@ -1458,6 +1468,61 @@ def test_resize(self):
  self.assertRaises(MemoryError, bytearray().resize, sys.maxsize)
  self.assertRaises(MemoryError, bytearray(1000).resize, sys.maxsize)
 
+ def test_take_bytes(self):
+ ba = bytearray(b'ab')
+ self.assertEqual(ba.take_bytes(), b'ab')
+ self.assertEqual(len(ba), 0)
+ self.assertEqual(ba, bytearray(b''))
+ self.assertEqual(ba.__alloc__(), 0)
+ base_size = sys.getsizeof(bytearray())
+ self.assertEqual(sys.getsizeof(ba), base_size)
+
+ # Positive and negative slicing.
+ ba = bytearray(b'abcdef')
+ self.assertEqual(ba.take_bytes(1), b'a')
+ self.assertEqual(ba, bytearray(b'bcdef'))
+ self.assertEqual(len(ba), 5)
+ self.assertEqual(ba.take_bytes(-5), b'')
+ self.assertEqual(ba, bytearray(b'bcdef'))
+ self.assertEqual(len(ba), 5)
+ self.assertEqual(ba.take_bytes(-3), b'bc')
+ self.assertEqual(ba, bytearray(b'def'))
+ self.assertEqual(len(ba), 3)
+ self.assertEqual(ba.take_bytes(3), b'def')
+ self.assertEqual(ba, bytearray(b''))
+ self.assertEqual(len(ba), 0)
+
+ # Take nothing from emptiness.
+ self.assertEqual(ba.take_bytes(0), b'')
+ self.assertEqual(ba.take_bytes(), b'')
+ self.assertEqual(ba.take_bytes(None), b'')
+
+ # Out of bounds, bad take value.
+ self.assertRaises(IndexError, ba.take_bytes, -1)
+ self.assertRaises(TypeError, ba.take_bytes, 3.14)
+ ba = bytearray(b'abcdef')
+ self.assertRaises(IndexError, ba.take_bytes, 7)
+
+ # Offset between physical and logical start (ob_bytes != ob_start).
+ ba = bytearray(b'abcde')
+ del ba[:2]
+ self.assertEqual(ba, bytearray(b'cde'))
+ self.assertEqual(ba.take_bytes(), b'cde')
+
+ # Overallocation at end.
+ ba = bytearray(b'abcde')
+ del ba[-2:]
+ self.assertEqual(ba, bytearray(b'abc'))
+ self.assertEqual(ba.take_bytes(), b'abc')
+ ba = bytearray(b'abcde')
+ ba.resize(4)
+ self.assertEqual(ba.take_bytes(), b'abcd')
+
+ # Take of a bytearray with references should fail.
+ ba = bytearray(b'abc')
+ with memoryview(ba) as mv:
+ self.assertRaises(BufferError, ba.take_bytes)
+ self.assertEqual(ba.take_bytes(), b'abc')
 
  def test_setitem(self):
  def setitem_as_mapping(b, i, val):
@@ -2564,6 +2629,18 @@ def zfill(b, a):
  c = a.zfill(0x400000)
  assert not c or c[-1] not in (0xdd, 0xcd)
 
+ def take_bytes(b, a): # MODIFIES!
+ b.wait()
+ c = a.take_bytes()
+ assert not c or c[0] == 48 # '0'
+
+ def take_bytes_n(b, a): # MODIFIES!
+ b.wait()
+ try:
+ c = a.take_bytes(10)
+ assert c == b'0123456789'
+ except IndexError: pass
+
  def check(funcs, a=None, *args):
  if a is None:
  a = bytearray(b'0' * 0x400000)
@@ -2625,6 +2702,10 @@ def check(funcs, a=None, *args):
  check([clear] + [startswith] * 10)
  check([clear] + [strip] * 10)
 
+ check([clear] + [take_bytes] * 10)
+ check([take_bytes_n] * 10, bytearray(b'0123456789' * 0x400))
+ check([take_bytes_n] * 10, bytearray(b'0123456789' * 5))
+
  check([clear] + [contains] * 10)
  check([clear] + [subscript] * 10)
  check([clear2] + [ass_subscript2] * 10, None, bytearray(b'0' * 0x400000))

Lib/test/test_capi/test_bytearray.py

@@ -1,3 +1,4 @@
+import sys
 import unittest
 from test.support import import_helper
 
@@ -55,7 +56,9 @@ def test_fromstringandsize(self):
  self.assertEqual(fromstringandsize(b'', 0), bytearray())
  self.assertEqual(fromstringandsize(NULL, 0), bytearray())
  self.assertEqual(len(fromstringandsize(NULL, 3)), 3)
- self.assertRaises(MemoryError, fromstringandsize, NULL, PY_SSIZE_T_MAX)
+ self.assertRaises(OverflowError, fromstringandsize, NULL, PY_SSIZE_T_MAX)
+ self.assertRaises(OverflowError, fromstringandsize, NULL,
+ PY_SSIZE_T_MAX-sys.getsizeof(b'') + 1)
 
  self.assertRaises(SystemError, fromstringandsize, b'abc', -1)
  self.assertRaises(SystemError, fromstringandsize, b'abc', PY_SSIZE_T_MIN)

Lib/test/test_sys.py

@@ -1583,7 +1583,7 @@ def test_objecttypes(self):
  samples = [b'', b'u'*100000]
  for sample in samples:
  x = bytearray(sample)
- check(x, vsize('n2Pi') + x.__alloc__())
+ check(x, vsize('n2PiP') + x.__alloc__())
  # bytearray_iterator
  check(iter(bytearray()), size('nP'))
  # bytes

Misc/NEWS.d/next/Core_and_Builtins/2025-10-14-18-24-16.gh-issue-139871.SWtuUz.rst

@@ -0,0 +1,2 @@
+Update :class:`bytearray` to use a :class:`bytes` under the hood as its buffer
+and add :func:`bytearray.take_bytes` to take it out.