from helper import unittest, PillowTestCase, tearDownModule from PIL import Image class TestPickle(PillowTestCase): def helper_pickle_file(self, pickle, protocol=0): im = Image.open('Tests/images/lena.jpg') filename = self.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 self.assertEqual(im, loaded_im) def helper_pickle_string( self, pickle, protocol=0, file='Tests/images/lena.jpg'): im = Image.open(file) # Act dumped_string = pickle.dumps(im, protocol) loaded_im = pickle.loads(dumped_string) # Assert self.assertEqual(im, loaded_im) def test_pickle_image(self): # Arrange import pickle # Act / Assert for protocol in range(0, pickle.HIGHEST_PROTOCOL + 1): self.helper_pickle_string(pickle, protocol) self.helper_pickle_file(pickle, protocol) def test_cpickle_image(self): # Arrange try: import cPickle except ImportError: return # Act / Assert for protocol in range(0, cPickle.HIGHEST_PROTOCOL + 1): self.helper_pickle_string(cPickle, protocol) self.helper_pickle_file(cPickle, protocol) def test_pickle_p_mode(self): # 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" ]: self.helper_pickle_string(pickle, file=file) if __name__ == '__main__': unittest.main() # End of file