Tests and partial implementation of pickling (#629)

This commit is contained in:
hugovk 2014-04-22 08:54:16 +03:00
parent 8c5ed8a873
commit adfbe8323a
2 changed files with 91 additions and 19 deletions

View File

@ -568,6 +568,15 @@ class Image:
return new return new
raise AttributeError(name) 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): def tobytes(self, encoder_name="raw", *args):
""" """
Return image as a bytes object Return image as a bytes object

63
Tests/test_pickle.py Normal file
View File

@ -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