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

Update `bytearray` to contain a `bytes` and provide a zero-copy path to
"extract" the `bytes`. This allows making several code paths more efficient.

This does not move any codepaths to make use of this new API. The documentation
changes include common code patterns which can be made more efficient with
this API.

---

When just changing `bytearray` to contain `bytes` I ran pyperformance on a
`--with-lto --enable-optimizations --with-static-libpython` build and don't see
any major speedups or slowdowns with this; all seems to be in the noise of
my machine (Generally changes under 5% or benchmarks that don't touch
bytes/bytearray).


Co-authored-by: Victor Stinner <vstinner@python.org>
Co-authored-by: Maurycy Pawłowski-Wieroński <5383+maurycy@users.noreply.github.com>

Author: 3 people

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.