From 2a6f2c5442d1095ed755d4fa699169f68bfceb1e Mon Sep 17 00:00:00 2001
From: hugovk <hugovk@users.noreply.github.com>
Date: Sat, 26 Apr 2014 19:43:53 +0300
Subject: [PATCH] Add __eq__ and __ne__ to Image to be able to test image
 equality when pickling. Pickle more data.

---
 PIL/Image.py         | 37 ++++++++++++++++++++++++++++++-------
 Tests/test_pickle.py | 33 +++++++++++++++++++--------------
 Tests/tester.py      |  3 ++-
 3 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/PIL/Image.py b/PIL/Image.py
index 889e92303..8d97e1221 100644
--- a/PIL/Image.py
+++ b/PIL/Image.py
@@ -563,6 +563,21 @@ class Image:
             self.save(file, format)
         return file
 
+    def __eq__(self, other):
+        a = (self.mode == other.mode)
+        b = (self.size == other.size)
+        c = (self.getpalette() == other.getpalette())
+        d = (self.info == other.info)
+        e = (self.category == other.category)
+        f = (self.readonly == other.readonly)
+        g = (self.pyaccess == other.pyaccess)
+        h = (self.tobytes() == other.tobytes())
+        return a and b and c and d and e and f and g and h
+
+    def __ne__(self, other):
+        eq = (self == other)
+        return not eq
+
     def __repr__(self):
         return "<%s.%s image mode=%s size=%dx%d at 0x%X>" % (
             self.__class__.__module__, self.__class__.__name__,
@@ -582,15 +597,23 @@ class Image:
         raise AttributeError(name)
 
     def __getstate__(self):
-        return [self.mode, self.size, self.tobytes()]
+        return [
+            self.info,
+            self.mode,
+            self.size,
+            self.getpalette(),
+            self.tobytes()]
 
     def __setstate__(self, state):
         Image.__init__(self)
         self.tile = []
-        mode, size, data = state
+        info, mode, size, palette, data = state
+        self.info = info
         self.mode = mode
         self.size = size
         self.im = core.new(mode, size)
+        if mode in ("L", "P"):
+            self.putpalette(palette)
         self.frombytes(data)
 
     def tobytes(self, encoder_name="raw", *args):
@@ -870,7 +893,7 @@ class Image:
 
         new_im = self._new(im)
         if delete_trns:
-            #crash fail if we leave a bytes transparency in an rgb/l mode.
+            # crash fail if we leave a bytes transparency in an rgb/l mode.
             del(new_im.info['transparency'])
         if trns is not None:
             if new_im.mode == 'P':
@@ -2188,8 +2211,8 @@ def open(fp, mode="r"):
                 fp.seek(0)
                 return factory(fp, filename)
         except (SyntaxError, IndexError, TypeError):
-            #import traceback
-            #traceback.print_exc()
+            # import traceback
+            # traceback.print_exc()
             pass
 
     if init():
@@ -2201,8 +2224,8 @@ def open(fp, mode="r"):
                     fp.seek(0)
                     return factory(fp, filename)
             except (SyntaxError, IndexError, TypeError):
-                #import traceback
-                #traceback.print_exc()
+                # import traceback
+                # traceback.print_exc()
                 pass
 
     raise IOError("cannot identify image file %r"
diff --git a/Tests/test_pickle.py b/Tests/test_pickle.py
index b52932757..7da385624 100644
--- a/Tests/test_pickle.py
+++ b/Tests/test_pickle.py
@@ -3,18 +3,6 @@ from tester import *
 from PIL import Image
 
 
-def test_frombytes_tobytes():
-    # Arrange
-    im = Image.open('Images/lena.jpg')
-
-    # Act
-    data = im.tobytes()
-    new_im = Image.frombytes(im.mode, im.size, data)
-
-    # Assert
-    assert_image_equal(im, new_im)
-
-
 def helper_test_pickle_file(pickle, protocol=0):
     im = Image.open('Images/lena.jpg')
     filename = tempfile('temp.pkl')
@@ -29,8 +17,8 @@ def helper_test_pickle_file(pickle, protocol=0):
     assert_image_equal(im, loaded_im)
 
 
-def helper_test_pickle_string(pickle, protocol=0):
-    im = Image.open('Images/lena.jpg')
+def helper_test_pickle_string(pickle, protocol=0, file='Images/lena.jpg'):
+    im = Image.open(file)
 
     # Act
     dumped_string = pickle.dumps(im, protocol)
@@ -62,4 +50,21 @@ def test_cpickle_image():
         helper_test_pickle_string(cPickle, protocol)
         helper_test_pickle_file(cPickle, protocol)
 
+
+def test_pickle_p_mode():
+    # Arrange
+    import pickle
+
+    # Act / Assert
+    for file in [
+            "Tests/images/test-card.png",
+            "Tests/images/zero_bb.png",
+            "Tests/images/zero_bb_scale2.png",
+            "Tests/images/non_zero_bb.png",
+            "Tests/images/non_zero_bb_scale2.png",
+            "Tests/images/p_trns_single.png",
+            "Tests/images/pil123p.png"
+    ]:
+        helper_test_pickle_string(pickle, file=file)
+
 # End of file
diff --git a/Tests/tester.py b/Tests/tester.py
index 32da48e98..a58872e2c 100644
--- a/Tests/tester.py
+++ b/Tests/tester.py
@@ -242,7 +242,8 @@ def assert_image_equal(a, b, msg=None):
         failure(msg or "got size %r, expected %r" % (a.size, b.size))
     elif a.tobytes() != b.tobytes():
         failure(msg or "got different content")
-        # generate better diff?
+    elif a != b:
+        failure(msg or "images different")
     else:
         success()