Add tp_richcompare handler for Imaging_Type/ImagingCore

Tests/test_lib_image.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import itertools
+
 import pytest
 
 from PIL import Image
@@ -32,3 +34,62 @@ def test_setmode() -> None:
         im.im.setmode("L")
     with pytest.raises(ValueError):
         im.im.setmode("RGBABCDE")
+
+
+@pytest.mark.parametrize("mode", Image.MODES)
+def test_equal(mode):
+    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.im == img_b.im
+
+
+# With mode "1" different bytes can map to the same value,
+# so we have to be more specific with the values we use.
+@pytest.mark.parametrize(
+    "bytes_a, bytes_b",
+    itertools.permutations(
+        (bytes(x) for x in itertools.product(b"\x00\xff", repeat=4)), 2
+    ),
+)
+def test_not_equal_mode_1(bytes_a, bytes_b):
+    # 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.im != img_b.im
+
+
+@pytest.mark.parametrize("mode", [mode for mode in Image.MODES if mode != "1"])
+def test_not_equal(mode):
+    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.im != img_b.im
+
+
+@pytest.mark.parametrize("mode", ("RGB", "YCbCr", "HSV", "LAB"))
+def test_equal_three_channels_four_bytes(mode):
+    # 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() == b"123"
+    assert img_b.tobytes() == b"123"
+    assert img_a.im == img_b.im
+
+
+@pytest.mark.parametrize("mode", ("LA", "La", "PA"))
+def test_equal_two_channels_four_bytes(mode):
+    img_a = Image.new(mode, (1, 1), 0x32000031)
+    img_b = Image.new(mode, (1, 1), 0x32FFFF31)
+    assert img_a.tobytes() == b"12"
+    assert img_b.tobytes() == b"12"
+    assert img_a.im == img_b.im

src/PIL/Image.py

@@ -668,13 +668,17 @@
         if self.__class__ is not other.__class__:
             return False
         assert isinstance(other, Image)
-        return (
+        if self is other:
+            return True
+        if not (
             self.mode == other.mode
             and self.size == other.size
             and self.info == other.info
-            and self.getpalette() == other.getpalette()
-            and self.tobytes() == other.tobytes()
-        )
+        ):
+            return False
+        self.load()
+        other.load()
+        return self.im == other.im
 
     def __repr__(self) -> str:
         return (

src/_imaging.c

@@ -3766,13 +3766,138 @@
     (ssizessizeobjargproc)NULL, /*sq_ass_slice*/
 };
 
+/*
+Returns 0 if all of the pixels are the same, otherwise 1.
+Skips unused bytes based on the given mode.
+*/
+static int
+_compare_pixels(
+    const char *mode,
+    const int ysize,
+    const int linesize,
+    const UINT8 **pixels_a,
+    const UINT8 **pixels_b
+) {
+    // Fortunately, all of the modes that have extra bytes in their pixels
+    // use four bytes for their pixels.
+    UINT32 mask = 0xffffffff;
+    if (!strcmp(mode, "RGB") || !strcmp(mode, "YCbCr") || !strcmp(mode, "HSV") ||
+        !strcmp(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
+    } else if (!strcmp(mode, "LA") || !strcmp(mode, "La") || !strcmp(mode, "PA")) {
+        // These modes have two channels in four bytes,
+        // so we have to ignore the middle two bytes.
+        mask = 0xff0000ff;
+    }
+
+    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 {
+        const int xsize = linesize / 4;
+        for (int y = 0; y < ysize; y++) {
+            UINT32 *line_a = (UINT32 *)pixels_a[y];
+            UINT32 *line_b = (UINT32 *)pixels_b[y];
+            for (int x = 0; x < xsize; x++, line_a++, line_b++) {
+                if ((*line_a & mask) != (*line_b & mask)) {
+                    return 1;
+                }
+            }
+        }
+    }
+    return 0;
+}
+
+static PyObject *
+image_richcompare(const ImagingObject *self, const PyObject *other, const int op) {
+    if (op != Py_EQ && op != Py_NE) {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+
+    // If the other object is not an ImagingObject.
+    if (!PyImaging_Check(other)) {
+        if (op == Py_EQ) {
+            Py_RETURN_FALSE;
+        } else {
+            Py_RETURN_TRUE;
+        }
+    }
+
+    const Imaging img_a = self->image;
+    const Imaging img_b = ((ImagingObject *)other)->image;
+
+    if (strcmp(img_a->mode, img_b->mode) || img_a->xsize != img_b->xsize ||
+        img_a->ysize != img_b->ysize) {
+        if (op == Py_EQ) {
+            Py_RETURN_FALSE;
+        } else {
+            Py_RETURN_TRUE;
+        }
+    }
+
+    const ImagingPalette palette_a = img_a->palette;
+    const ImagingPalette palette_b = img_b->palette;
+    if (palette_a || palette_b) {
+        const UINT8 *palette_a_data = palette_a->palette;
+        const UINT8 *palette_b_data = palette_b->palette;
+        const UINT8 **palette_a_data_ptr = &palette_a_data;
+        const UINT8 **palette_b_data_ptr = &palette_b_data;
+        if (!palette_a || !palette_b || palette_a->size != palette_b->size ||
+            strcmp(palette_a->mode, palette_b->mode) ||
+            _compare_pixels(
+                palette_a->mode,
+                1,
+                palette_a->size * 4,
+                palette_a_data_ptr,
+                palette_b_data_ptr
+            )) {
+            if (op == Py_EQ) {
+                Py_RETURN_FALSE;
+            } else {
+                Py_RETURN_TRUE;
+            }
+        }
+    }
+
+    if (_compare_pixels(
+            img_a->mode,
+            img_a->ysize,
+            img_a->linesize,
+            (const UINT8 **)img_a->image,
+            (const UINT8 **)img_b->image
+        )) {
+        if (op == Py_EQ) {
+            Py_RETURN_FALSE;
+        } else {
+            Py_RETURN_TRUE;
+        }
+    } else {
+        if (op == Py_EQ) {
+            Py_RETURN_TRUE;
+        } else {
+            Py_RETURN_FALSE;
+        }
+    }
+}
+
 /* 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 = (richcmpfunc)image_richcompare,
     .tp_methods = methods,
     .tp_getset = getsetters,
 };