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.