Speed up image comparison and reduce memory usage

Tests/test_image_equality.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import itertools
+from unittest.mock import Mock
+
+import pytest
+
+from PIL import Image
+
+
+def _make_rows(width: int, height: int) -> list[bytes]:
+    """Build test image data of width x height."""
+    return [
+        bytes((y * (width - 1) + x) % 256 for x in range(width)) for y in range(height)
+    ]
+
+
+def _frombuffer(
+    mode: str, size: tuple[int, int], data: bytes, *, stride: int = 0, ystep: int = 1
+) -> Image.Image:
+    im = Image.frombuffer(mode, size, data, "raw", mode, stride, ystep)
+    assert im.readonly  # Sanity check (that we took the map_buffer path)
+    return im
+
+
+@pytest.mark.parametrize("mode", Image.MODES)
+def test_equal(mode: str) -> None:
+    num_img_bytes = len(Image.new(mode, (2, 2)).tobytes())
+    data = bytes(range(ord("A"), ord("A") + num_img_bytes))
+    img_a = Image.frombytes(mode, (2, 2), data)
+    img_b = Image.frombytes(mode, (2, 2), data)
+    assert img_a.tobytes() == img_b.tobytes()
+    assert img_a == img_b
+
+
+def test_not_equal_mode_1() -> None:
+    # With mode "1" different bytes can map to the same value,
+    # so we have to be more specific with the values we use.
+    for bytes_a, bytes_b in itertools.permutations(
+        (bytes(x) for x in itertools.product(b"\x00\xff", repeat=4)), 2
+    ):
+        # Use rawmode "1;8" so that each full byte is interpreted as a value
+        # instead of the bits in the bytes being interpreted as values.
+        img_a = Image.frombytes("1", (2, 2), bytes_a, "raw", "1;8")
+        img_b = Image.frombytes("1", (2, 2), bytes_b, "raw", "1;8")
+        assert img_a.tobytes() != img_b.tobytes()
+        assert img_a != img_b
+
+
+@pytest.mark.parametrize("mode", [mode for mode in Image.MODES if mode != "1"])
+def test_not_equal(mode: str) -> None:
+    num_img_bytes = len(Image.new(mode, (2, 2)).tobytes())
+    data_a = bytes(range(ord("A"), ord("A") + num_img_bytes))
+    data_b = bytes(range(ord("Z"), ord("Z") - num_img_bytes, -1))
+    img_a = Image.frombytes(mode, (2, 2), data_a)
+    img_b = Image.frombytes(mode, (2, 2), data_b)
+    assert img_a.tobytes() != img_b.tobytes()
+    assert img_a != img_b
+
+
+@pytest.mark.parametrize("mode", ("RGB", "YCbCr", "HSV", "LAB"))
+def test_equal_three_channels_four_bytes(mode: str) -> None:
+    # The "A" and "B" values in LAB images are signed values from -128 to 127,
+    # but we store them as unsigned values from 0 to 255, so we need to use
+    # slightly different input bytes for LAB to get the same output.
+    img_a = Image.new(mode, (1, 1), 0x00B3B231 if mode == "LAB" else 0x00333231)
+    img_b = Image.new(mode, (1, 1), 0xFFB3B231 if mode == "LAB" else 0xFF333231)
+    assert img_a.tobytes() == img_b.tobytes() == b"123"
+    assert img_a == img_b
+
+
+@pytest.mark.parametrize("mode", ("LA", "La", "PA"))
+def test_equal_two_channels_four_bytes(mode: str) -> None:
+    # Test that for LA/La/PA modes, where the data is stored in the 1st and 4th
+    # byte of each pixel, the middle bytes are masked off for comparison.
+    img_a = Image.new(mode, (1, 1), 0x32000031)
+    img_b = Image.new(mode, (1, 1), 0x32FFFF31)
+    assert img_a.tobytes() == img_b.tobytes() == b"12"
+    assert img_a == img_b
+
+
+def test_not_equal_rgbx_padding() -> None:
+    # Ensure RGBX's last byte is compared, even if it has no image meaning.
+    img_a = Image.frombytes("RGBX", (1, 1), b"1234")
+    img_b = Image.frombytes("RGBX", (1, 1), b"123\xff")
+    assert img_a.tobytes() != img_b.tobytes()
+    assert img_a != img_b
+
+
+def test_compare_with_other_type() -> None:
+    im = Image.new("L", (1, 1))
+    assert im.im == im.im
+    assert im.im != 42
+    # Check that the other object's comparison method is called.
+    x = Mock(__eq__=Mock(return_value=True))
+    assert im.im == x
+    assert x.__eq__.called  # type: ignore[attr-defined]
+
+
+def test_equal_frombuffer_stride() -> None:
+    # Test that a buffer-mapped image's padding bytes (stride)
+    # are not part of the comparison.
+    width, height, stride = 3, 4, 8
+    rows = _make_rows(width, height)
+    buffer_a = b"".join(row + b"\x00" * (stride - width) for row in rows)
+    buffer_b = b"".join(row + b"\xff" * (stride - width) for row in rows)
+    assert buffer_a != buffer_b
+    img_a = _frombuffer("L", (width, height), buffer_a, stride=stride)
+    img_b = _frombuffer("L", (width, height), buffer_b, stride=stride)
+    assert img_a.tobytes() == img_b.tobytes()  # Padding bytes disappear
+    assert img_a == img_b
+
+
+def test_not_equal_frombuffer_stride() -> None:
+    # Test that differences within strided rows are found,
+    # even if padding matches.
+    width, height, stride = 3, 4, 8
+    rows = _make_rows(width, height)
+    padding = b"\xab" * (stride - width)
+    buffer_a = b"".join(row + padding for row in rows)
+    rows[height - 1] = bytes(42) + rows[height - 1][1:]
+    buffer_b = b"".join(row + padding for row in rows)
+
+    img_a = _frombuffer("L", (width, height), buffer_a, stride=stride)
+    img_b = _frombuffer("L", (width, height), buffer_b, stride=stride)
+    assert img_a.tobytes() != img_b.tobytes()
+    assert img_a != img_b
+
+
+def test_equal_frombuffer_ystep() -> None:
+    # Test that ystep=-1 (rows in reverse order) is handled correctly.
+    width, height = 3, 4
+    rows = _make_rows(width, height)
+    img_a = _frombuffer("L", (width, height), b"".join(rows))
+    img_b = _frombuffer("L", (width, height), b"".join(reversed(rows)), ystep=-1)
+    assert img_a.tobytes() == img_b.tobytes()
+    assert img_a == img_b

docs/releasenotes/13.0.0.rst

@@ -228,3 +228,8 @@ position or contents.
 
 :py:meth:`~PIL.Image.Image.tobytes` also now explicitly cleans up its encoder on
 both success and failure, rather than relying on object destruction.
+
+Faster image comparison
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Comparing two images with ``==`` was made faster by reducing data copying.

src/PIL/Image.py

@@ -754,13 +754,19 @@ def __eq__(self, other: object) -> bool:
         if self.__class__ is not other.__class__:
             return False
         assert isinstance(other, Image)
-        return (
-            self.mode == other.mode
-            and self.size == other.size
-            and self.info == other.info
-            and self.getpalette() == other.getpalette()
-            and self.tobytes() == other.tobytes()
-        )
+        if self is other:
+            return True
+        # Early-out before loading.
+        # ImagingCore does check for mode and size.
+        if (
+            self.mode != other.mode
+            or self.size != other.size
+            or self.info != other.info
+        ):
+            return False
+        self.load()
+        other.load()
+        return self.im == other.im
 
     def __repr__(self) -> str:
         return (

src/_imaging.c

@@ -3883,13 +3883,168 @@ static PySequenceMethods image_as_sequence = {
     (ssizessizeobjargproc)NULL, /*sq_ass_slice*/
 };
 
+/**
+ * Compare two pixel arrays in the given mode for equality.
+ *
+ * @param mode The mode of the pixel arrays.
+ * @param ysize The number of rows in the pixel arrays.
+ * @param linesize The number of bytes in each row of the pixel arrays.
+ * @param pixels_a Pointer to the first array of pixel row arrays.
+ * @param pixels_b Pointer to the second array of pixel row arrays.
+ * @return 0 if the pixel arrays are equal, 1 if they are not equal.
+ */
+static int
+_compare_pixels(
+    const ModeID mode,
+    const int ysize,
+    const int linesize,
+    const char *const *pixels_a,
+    const char *const *pixels_b
+) {
+    if (ysize <= 0 || linesize <= 0) {
+        // Pixel arrays without size are always equal.
+        // Since rows may not have been allocated at all,
+        // early-out now.
+        return 0;
+    }
+
+    // Fortunately, all of the modes that have extra bytes in their pixels
+    // use four bytes for their pixels. RGBX is deliberately not one of them:
+    // its fourth byte is compared just like Image.tobytes() includes it.
+    UINT32 mask = 0xffffffff;
+    switch (mode) {
+        case IMAGING_MODE_RGB:
+        case IMAGING_MODE_YCbCr:
+        case IMAGING_MODE_HSV:
+        case IMAGING_MODE_LAB:
+            // These modes have three channels in four bytes,
+            // so we have to ignore the last byte.
+#ifdef WORDS_BIGENDIAN
+            mask = 0xffffff00;
+#else
+            mask = 0x00ffffff;
+#endif
+            break;
+        case IMAGING_MODE_LA:
+        case IMAGING_MODE_La:
+        case IMAGING_MODE_PA:
+            // These modes have two channels in four bytes,
+            // so we have to ignore the middle two bytes.
+            mask = 0xff0000ff;
+            break;
+        default:
+            break;
+    }
+
+    if (mask == 0xffffffff) {
+        // If we aren't masking anything we can use memcmp.
+        for (int y = 0; y < ysize; y++) {
+            if (memcmp(pixels_a[y], pixels_b[y], linesize)) {
+                return 1;
+            }
+        }
+    } else {
+        // All modes where we use a mask store four bytes per pixel.
+        const int xsize = linesize / 4;
+        for (int y = 0; y < ysize; y++) {
+            const UINT32 *line_a = (const UINT32 *)pixels_a[y];
+            const UINT32 *line_b = (const UINT32 *)pixels_b[y];
+            // Check the first pixel separately; if it's different,
+            // we don't need to enter the loop at all.
+            if ((line_a[0] ^ line_b[0]) & mask) {
+                return 1;
+            }
+            // Autovectorizable form: compute one row at a time,
+            // and OR together all the differences in that row,
+            // then see if we saw any differences at all.
+            // This loop will not bail out in the middle of a row
+            // (the common case of "entirely different images,
+            // same size and mode" will have been taken care of
+            // above in the first-pixel check), but it measures
+            // to be faster anyhow on modern CPUs.
+            UINT32 diff = 0;
+            for (int x = 0; x < xsize; x++) {
+                diff |= (line_a[x] ^ line_b[x]) & mask;
+            }
+            if (diff) {
+                return 1;
+            }
+        }
+    }
+    return 0;
+}
+
+/**
+ * Compare two palettes for equality.
+ *
+ * @param palette_a Pointer to the first palette. May be NULL.
+ * @param palette_b Pointer to the second palette. May be NULL.
+ * @return 0 if the palettes are equal, 1 if they are not equal.
+ */
+static int
+_compare_palette(const ImagingPalette palette_a, const ImagingPalette palette_b) {
+    if (palette_a == NULL && palette_b == NULL) {
+        return 0;
+    }
+    if (palette_a == NULL || palette_b == NULL) {
+        return 1;
+    }
+    if (palette_a->size != palette_b->size || palette_a->mode != palette_b->mode) {
+        return 1;
+    }
+
+    const char *palette_a_data = (const char *)palette_a->palette;
+    const char *palette_b_data = (const char *)palette_b->palette;
+    // 1024 is the hard-coded maximum size of a palette
+    int linesize = MIN(palette_a->size * 4, 1024);
+    return _compare_pixels(
+        palette_a->mode, 1, linesize, &palette_a_data, &palette_b_data
+    );
+}
+
+static PyObject *
+image_richcompare(PyObject *self, PyObject *other, int op) {
+    if (op != Py_EQ && op != Py_NE) {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+
+    if (!PyImaging_Check(other)) {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+
+    const Imaging img_a = ((ImagingObject *)self)->image;
+    const Imaging img_b = ((ImagingObject *)other)->image;
+
+    if (img_a == NULL || img_b == NULL) {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+
+    const int equal =
+        // Mode and size comparison (cheap)
+        img_a->mode == img_b->mode && img_a->xsize == img_b->xsize &&
+        img_a->ysize == img_b->ysize &&
+        // Palette comparison
+        _compare_palette(img_a->palette, img_b->palette) == 0 &&
+        // Pixel comparison
+        _compare_pixels(
+            img_a->mode,
+            img_a->ysize,
+            img_a->linesize,
+            (const char *const *)img_a->image,
+            (const char *const *)img_b->image
+        ) == 0;
+
+    return PyBool_FromLong(op == Py_EQ ? equal : !equal);
+}
+
 /* type description */
 
 static PyTypeObject Imaging_Type = {
     PyVarObject_HEAD_INIT(NULL, 0).tp_name = "ImagingCore",
     .tp_basicsize = sizeof(ImagingObject),
     .tp_dealloc = (destructor)_dealloc,
     .tp_as_sequence = &image_as_sequence,
+    .tp_richcompare = image_richcompare,
     .tp_methods = methods,
     .tp_getset = getsetters,
 };