Improve memory usage and performance for tobytes()

Tests/test_image_encode.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import pytest
+
+from PIL import Image
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@pytest.mark.parametrize(
+    "mode",
+    ("1", "L", "P", "LA", "I", "F", "I;16", "RGB", "RGBA", "CMYK"),
+)
+def test_raw_encoder_optimal_bufsize(mode: str) -> None:
+    # Test the raw encoder does know the exact buffer size.
+    im = Image.new(mode, (300, 200))
+    encoder = Image._getencoder(mode, "raw", mode)
+    encoder.setimage(im.im, (0, 0) + im.size)
+    assert encoder.optimal_bufsize == len(im.tobytes())
+
+
+def test_raw_encoder_optimal_bufsize_stride() -> None:
+    im = Image.new("RGB", (100, 50))
+    encoder = Image._getencoder("RGB", "raw", ("RGB", 400))
+    encoder.setimage(im.im, (0, 0) + im.size)
+    assert encoder.optimal_bufsize == 400 * 50
+    encoder.encode(400)  # The encoder's internals change here
+    # ... but the result for this should remain the same.
+    assert encoder.optimal_bufsize == 400 * 50
+
+
+def test_encoder_optimal_bufsize_unknown() -> None:
+    im = Image.new("1", (64, 64))
+
+    # the size is not known before the image has been assigned
+    encoder = Image._getencoder("1", "raw", "1")
+    assert encoder.optimal_bufsize == 0
+
+    # nor for encoders whose output size depends on the pixel data
+    encoder = Image._getencoder("1", "xbm", "1")
+    encoder.setimage(im.im, (0, 0) + im.size)
+    assert encoder.optimal_bufsize == 0
+
+
+class EncoderSpy:
+    def __init__(self, encoder: Any, *, bufsizes: list[int]) -> None:
+        self._encoder = encoder
+        self.bufsizes = bufsizes
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._encoder, name)
+
+    def encode(self, bufsize: int) -> tuple[int, int, bytes]:
+        self.bufsizes.append(bufsize)
+        return self._encoder.encode(bufsize)
+
+
+@pytest.mark.parametrize("mode", ("L", "RGB"))
+@pytest.mark.parametrize(
+    "size",
+    (
+        (20000, 1),
+        (1, 20000),
+        (3, 3),
+        (500, 500),
+    ),
+)
+def test_tobytes_exact_buffer_and_single_pass(
+    mode: str,
+    size: tuple[int, int],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """
+    Test that the raw encoder gets asked exactly the correct size,
+    and ends up doing its work in a single pass.
+    """
+    bufsizes: list[int] = []
+    get_encoder = Image._getencoder
+    monkeypatch.setattr(
+        Image,
+        "_getencoder",
+        lambda *args: EncoderSpy(get_encoder(*args), bufsizes=bufsizes),
+    )
+
+    im = Image.new(mode, size)
+    n_bands = len(im.getbands())
+    expected_size = size[0] * size[1] * n_bands
+    assert len(im.tobytes()) == expected_size
+    assert bufsizes == [expected_size]

docs/releasenotes/13.0.0.rst

@@ -104,6 +104,33 @@ Two new filters are available for :py:meth:`~PIL.Image.Image.resize` and
 Other changes
 =============
 
+Reduced memory use in Image.tobytes() in raw mode
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Encoders may now report the "optimal" buffer size that they will need to encode the
+entire image into their output format.
+
+In general, :py:meth:`~PIL.Image.Image.tobytes` collects encoded data in
+``ImageFile.MAXBLOCK`` sized chunks and join them into the output byte buffer
+afterwards. For raw uncompressed output, this would need, momentarily, roughly
+twice the memory of the returned data. The raw encoder now reports the buffer size
+that encodes the whole image in a single call, so no copy of the result is made.
+
+Encoders that cannot know their output size in advance, such as the compressed
+encoders, are unaffected and keep using :py:data:`~PIL.ImageFile.MAXBLOCK`;
+for many applications, it may be beneficial to adjust that value for improved
+performance and reduced peak memory use.
+
+This also reduces peak memory use in other operations in Pillow that use this
+path for encoding, including but not limited to:
+
+* AVIF and SGI encoding
+* BLP JPEG decoding
+* Image equivalence comparisons
+* NumPy array interface support
+* Pickling of images
+* Qt interop
+
 Python 3.15
 ^^^^^^^^^^^
 

src/PIL/Image.py

@@ -886,7 +886,24 @@ def tobytes(self, encoder_name: str = "raw", *args: Any) -> bytes:
 
         from . import ImageFile
 
-        bufsize = max(ImageFile.MAXBLOCK, self.size[0] * 4)  # see RawEncode.c
+        # Figure out an optimal buffer size. A buffer size that's too small
+        # to hold all of a given image's encoded bytes will result in `output`
+        # having a bunch of small chunks, and returning the actual bytes for
+        # them will then require a second contiguous allocation for the result,
+        # leading to double-plus-then-some peak memory allocation.
+        #
+        # CPython optimizes the case `b"".join([x])` for a simple `x` to be
+        # a reference increase of `x` and a return instead of any extra allocations;
+        # this has been the case since approximately 2004 (commit 05eba1fdc80).
+        #
+        # * e.optimal_bufsize will be nonzero for encoders that can compute it.
+        # * The heuristic `self.size[0] * 4` for the buffer size matches the widest
+        #   possible row of pixels in a RGBA image; conveniently, this is also slightly
+        #   larger than what is required for a single row in XbmEncode.c,
+        #   which in itself is unable to cope with a too-small buffer.
+        # * MAXBLOCK can be tuned by the user; for the aforementioned reasons,
+        #   we need to use the larger of that and the heuristic.
+        bufsize = e.optimal_bufsize or max(ImageFile.MAXBLOCK, self.size[0] * 4)
 
         output = []
         while True:

src/PIL/ImageFile.py

@@ -892,6 +892,14 @@ class PyEncoder(PyCodec):
     def pushes_fd(self) -> bool:
         return self._pushes_fd
 
+    @property
+    def optimal_bufsize(self) -> int:
+        """
+        The buffer size that would encode the whole tile in a single
+        :meth:`encode` call, or 0 if that is not known in advance.
+        """
+        return 0
+
     def encode(self, bufsize: int) -> tuple[int, int, bytes]:
         """
         Override to perform the encoding process.

src/encode.c

@@ -43,6 +43,7 @@ typedef struct {
         Imaging im, ImagingCodecState state, UINT8 *buffer, int bytes
     );
     int (*cleanup)(ImagingCodecState state);
+    Py_ssize_t (*optimal_bufsize)(ImagingCodecState state);
     struct ImagingCodecStateInstance state;
     Imaging im;
     PyObject *lock;
@@ -83,8 +84,9 @@ PyImaging_EncoderNew(int contextsize) {
     /* Initialize encoder context */
     encoder->state.context = context;
 
-    /* Most encoders don't need this */
+    /* Most encoders don't need these */
     encoder->cleanup = NULL;
+    encoder->optimal_bufsize = NULL;
 
     /* Target image */
     encoder->lock = NULL;
@@ -337,12 +339,31 @@ static struct PyMethodDef methods[] = {
     {NULL, NULL} /* sentinel */
 };
 
+/* The buffer size that would let this encoder consume the whole tile in a
+   single call, or 0 if that size cannot be known up front (basically all
+   non-raw formats). Only meaningful once setimage() has been called. */
+static PyObject *
+_get_optimal_bufsize(ImagingEncoderObject *encoder, void *closure) {
+    Py_ssize_t bufsize = 0;
+
+    if (encoder->optimal_bufsize && encoder->im) {
+        bufsize = encoder->optimal_bufsize(&encoder->state);
+    }
+
+    return PyLong_FromSsize_t(bufsize);
+}
+
 static struct PyGetSetDef getseters[] = {
     {"pushes_fd",
      (getter)_get_pushes_fd,
      NULL,
      "True if this decoder expects to push directly to self.fd",
      NULL},
+    {"optimal_bufsize",
+     (getter)_get_optimal_bufsize,
+     NULL,
+     "Buffer size that encodes the entire tile in one call, or 0 if unknown",
+     NULL},
     {NULL, NULL, NULL, NULL, NULL} /* sentinel */
 };
 
@@ -497,6 +518,40 @@ PyImaging_PcxEncoderNew(PyObject *self, PyObject *args) {
 /* RAW                                                                  */
 /* -------------------------------------------------------------------- */
 
+static Py_ssize_t
+_raw_optimal_bufsize(ImagingCodecState state) {
+    // ImagingRawEncode writes one row at a time into a buffer of at least `bytes`.
+    // Before the first `encode()` call, `bytes` is the packed row size, and `count`
+    // is the stride (or zero). After the first call, the two get swapped...
+    // The larger is the row size.
+    Py_ssize_t row = MAX(state->bytes, state->count);
+    Py_ssize_t bufsize = row * state->ysize;
+    // Cap the size to the whole number of rows that fits an `int`.
+    // TODO: this needs to be changed when encoders' buffer sizes become `ssize_t`,
+    //       like decoders did in ca1cf5925.
+    if (bufsize > INT_MAX) {
+        bufsize = INT_MAX - (INT_MAX % row);
+    }
+    return bufsize;
+}
+
+/**
+ * Instantiate a raw, uncompressed encoder.
+ *
+ * Python arguments:
+ * - mode (str): The mode name for finding a pixel packer.
+ * - rawmode (str): The rawmode name for finding a pixel packer.
+ * - stride (int, optional): The number of bytes each row occupies in the output stream,
+ *                           >= the packed pixel size. Each row is explicitly
+ *                           zero-padded to this size. If unset, the true packed size
+ *                           is used.
+ * - ystep (int, optional): If set to a negative value, the encoder will write rows in
+ *                          reverse order. The only effective values are -1 and +1;
+ *                          this does not have the encoder skip rows.
+ * @param self Unused.
+ * @param args Python arguments, see above.
+ * @return A Python object representing an encoder.
+ */
 PyObject *
 PyImaging_RawEncoderNew(PyObject *self, PyObject *args) {
     ImagingEncoderObject *encoder;
@@ -523,6 +578,7 @@ PyImaging_RawEncoderNew(PyObject *self, PyObject *args) {
     }
 
     encoder->encode = ImagingRawEncode;
+    encoder->optimal_bufsize = _raw_optimal_bufsize;
 
     encoder->state.ystep = ystep;
     encoder->state.count = stride;