Tests and partial implementation of pickling (#629)

PIL/Image.py

@@ -101,7 +101,7 @@ import collections
 import numbers
 
 # works everywhere, win for pypy, not cpython
-USE_CFFI_ACCESS = hasattr(sys, 'pypy_version_info')  
+USE_CFFI_ACCESS = hasattr(sys, 'pypy_version_info')
 try:
     import cffi
     HAS_CFFI=True
@@ -233,7 +233,7 @@ _MODE_CONV = {
     "CMYK": ('|u1', 4),
     "YCbCr": ('|u1', 3),
     "LAB": ('|u1', 3), # UNDONE - unsigned |u1i1i1
-	# I;16 == I;16L, and I;32 == I;32L  
+	# I;16 == I;16L, and I;32 == I;32L
     "I;16": ('<u2', None),
     "I;16B": ('>u2', None),
     "I;16L": ('<u2', None),
@@ -502,7 +502,7 @@ class Image:
         return self
     def __exit__(self, *args):
         self.close()
-        
+
     def close(self):
         """
         Closes the file pointer, if possible.
@@ -524,7 +524,7 @@ class Image:
         # deferred error that will better explain that the core image
         # object is gone.
         self.im = deferred_error(ValueError("Operation on closed image"))
-        
+
 
     def _copy(self):
         self.load()
@@ -540,7 +540,7 @@ class Image:
         if not file:
             f, file = tempfile.mkstemp(suffix)
             os.close(f)
-            
+
         self.load()
         if not format or format == "PPM":
             self.im.save_ppm(file)
@@ -568,6 +568,15 @@ class Image:
             return new
         raise AttributeError(name)
 
+    def __getstate__(self):
+        return [self.mode, self.size, self.tobytes()]
+
+    def __setstate__(self, state):
+        mode, size, data = state
+        self.mode = mode
+        self.size = size
+        self.frombytes(data)
+
     def tobytes(self, encoder_name="raw", *args):
         """
         Return image as a bytes object
@@ -672,7 +681,7 @@ class Image:
         normal cases, you don't need to call this method, since the
         Image class automatically loads an opened image when it is
         accessed for the first time. This method will close the file
-        associated with the image. 
+        associated with the image.
 
         :returns: An image access object.
         """
@@ -777,7 +786,7 @@ class Image:
         if "transparency" in self.info and self.info['transparency'] is not None:
             if self.mode in ('L', 'RGB') and mode == 'RGBA':
                 # Use transparent conversion to promote from transparent
-                # color to an alpha channel.         
+                # color to an alpha channel.
                 return self._new(self.im.convert_transparent(
                                            mode, self.info['transparency']))
             elif self.mode in ('L', 'RGB', 'P') and mode in ('L', 'RGB', 'P'):
@@ -799,11 +808,11 @@ class Image:
                         trns_im = trns_im.convert(mode)
                     else:
                         # can't just retrieve the palette number, got to do it
-                        # after quantization. 
+                        # after quantization.
                         trns_im = trns_im.convert('RGB')
                     trns = trns_im.getpixel((0,0))
-                        
+
-                    
+
         if mode == "P" and palette == ADAPTIVE:
             im = self.im.quantize(colors)
             new = self._new(im)
@@ -811,7 +820,7 @@ class Image:
             new.palette = ImagePalette.raw("RGB", new.im.getpalette("RGB"))
             if delete_trns:
                 # This could possibly happen if we requantize to fewer colors.
-                # The transparency would be totally off in that case. 
+                # The transparency would be totally off in that case.
                 del(new.info['transparency'])
             if trns is not None:
                 try:
@@ -826,7 +835,7 @@ class Image:
         # colorspace conversion
         if dither is None:
             dither = FLOYDSTEINBERG
-                
+
         try:
             im = self.im.convert(mode, dither)
         except ValueError:
@@ -863,7 +872,7 @@ class Image:
         # quantizer interface in a later version of PIL.
 
         self.load()
-        
+
         if method is None:
             # defaults:
             method = 0
@@ -871,10 +880,10 @@ class Image:
                 method = 2
 
         if self.mode == 'RGBA' and method != 2:
-            # Caller specified an invalid mode. 
+            # Caller specified an invalid mode.
             raise ValueError('Fast Octree (method == 2) is the ' +
                              ' only valid method for quantizing RGBA images')
-        
+
         if palette:
             # use palette from reference image
             palette.load()
@@ -928,7 +937,7 @@ class Image:
     def draft(self, mode, size):
         """
         NYI
-        
+
         Configures the image file loader so it returns a version of the
         image that as closely as possible matches the given mode and
         size.  For example, you can use this method to convert a color
@@ -1277,7 +1286,7 @@ class Image:
             if self.mode in ("I", "I;16", "F"):
                 # check if the function can be used with point_transform
                 # UNDONE wiredfool -- I think this prevents us from ever doing
-                # a gamma function point transform on > 8bit images. 
+                # a gamma function point transform on > 8bit images.
                 scale, offset = _getscaleoffset(lut)
                 return self._new(self.im.point_transform(scale, offset))
             # for other modes, convert the function to a table
@@ -1420,8 +1429,8 @@ class Image:
             self._copy()
             self.pyaccess = None
             self.load()
-            
+
-        if self.pyaccess: 
+        if self.pyaccess:
             return self.pyaccess.putpixel(xy,value)
         return self.im.putpixel(xy, value)
 

Tests/test_pickle.py

@@ -0,0 +1,63 @@
+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')
+
+    # Act
+    with open(filename, 'wb') as f:
+        pickle.dump(im, f, protocol)
+    with open(filename, 'rb') as f:
+        loaded_im = pickle.load(f)
+
+    # Assert
+    assert_image_equal(im, loaded_im)
+
+
+def helper_test_pickle_string(pickle, protocol=0):
+    im = Image.open('Images/lena.jpg')
+
+    # Act
+    dumped_string = pickle.dumps(im, protocol)
+    loaded_im = pickle.loads(dumped_string)
+
+    # Assert
+    assert_image_equal(im, loaded_im)
+
+
+def test_pickle_image():
+    # Arrange
+    import pickle
+
+    # Act / Assert
+    for protocol in range(0, pickle.HIGHEST_PROTOCOL + 1):
+        helper_test_pickle_string(pickle, protocol)
+        helper_test_pickle_file(pickle, protocol)
+
+
+def test_cpickle_image():
+    # Arrange
+    import cPickle
+
+    # Act / Assert
+    for protocol in range(0, cPickle.HIGHEST_PROTOCOL + 1):
+        helper_test_pickle_string(cPickle, protocol)
+        helper_test_pickle_file(cPickle, protocol)
+
+
+# End of file