diff -r 92003d9aae0e Lib/test/test_pickle.py --- a/Lib/test/test_pickle.py Tue Feb 26 12:37:07 2013 +0000 +++ b/Lib/test/test_pickle.py Wed Feb 27 22:42:21 2013 +0530 @@ -1,5 +1,7 @@ import pickle -from cStringIO import StringIO +import cStringIO +import io +import os from test import test_support @@ -22,46 +24,95 @@ module = pickle error = KeyError -class PicklerTests(AbstractPickleTests): +class AbstractIOPicklerTests(AbstractPickleTests): error = KeyError + fileFlag = 0 + + def teardown(self): + if(self.f): + self.f.close() + os.remove(test_support.TESTFN) def dumps(self, arg, proto=0, fast=0): - f = StringIO() - p = pickle.Pickler(f, proto) + if not self.fileFlag: + self.f = self.ioclass() + else: + self.f = self.ioclass(file=test_support.TESTFN, mode='w+') + p = pickle.Pickler(self.f, proto) if fast: p.fast = fast p.dump(arg) - f.seek(0) - return f.read() + self.f.seek(0) + return self.f.read() def loads(self, buf): - f = StringIO(buf) - u = pickle.Unpickler(f) + if not self.fileFlag: + self.f = self.ioclass(buf) + else: + self.f = self.ioclass(file=test_support.TESTFN, mode='w+') + self.f.write(buf) + self.f.seek(0) + u = pickle.Unpickler(self.f) return u.load() -class PersPicklerTests(AbstractPersistentPicklerTests): +class StringIOPicklerTests(AbstractIOPicklerTests): + ioclass = cStringIO.StringIO + +class BytesIOPicklerTests(AbstractIOPicklerTests): + ioclass = io.BytesIO + +class FileIOPicklerTests(AbstractIOPicklerTests): + ioclass = io.FileIO + fileFlag = 1 + +class AbstractIOPersPicklerTests(AbstractPersistentPicklerTests): + + fileFlag = 0 + + def teardown(self): + if(self.f): + self.f.close() + os.remove(test_support.TESTFN) def dumps(self, arg, proto=0, fast=0): class PersPickler(pickle.Pickler): def persistent_id(subself, obj): return self.persistent_id(obj) - f = StringIO() - p = PersPickler(f, proto) + if not self.fileFlag: + self.f = self.ioclass() + else: + self.f = self.ioclass(file=test_support.TESTFN, mode='w+') + p = PersPickler(self.f, proto) if fast: p.fast = fast p.dump(arg) - f.seek(0) - return f.read() + self.f.seek(0) + return self.f.read() def loads(self, buf): class PersUnpickler(pickle.Unpickler): def persistent_load(subself, obj): return self.persistent_load(obj) - f = StringIO(buf) - u = PersUnpickler(f) + if not self.fileFlag: + self.f = self.ioclass(buf) + else: + self.f = self.ioclass(file=test_support.TESTFN, mode='w+') + self.f.write(buf) + self.f.seek(0) + u = PersUnpickler(self.f) return u.load() +class StringIOPersPicklerTests(AbstractIOPersPicklerTests): + ioclass = cStringIO.StringIO + +class BytesIOPersPicklerTests(AbstractIOPersPicklerTests): + ioclass = io.BytesIO + +class FileIOPersPicklerTests(AbstractIOPersPicklerTests): + ioclass = io.FileIO + fileFlag = 1 + class PicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests): pickler_class = pickle.Pickler @@ -81,12 +132,18 @@ def test_main(): test_support.run_unittest( PickleTests, - PicklerTests, - PersPicklerTests, + StringIOPicklerTests, + BytesIOPicklerTests, + FileIOPicklerTests, + StringIOPersPicklerTests, + BytesIOPersPicklerTests, + FileIOPersPicklerTests, PicklerUnpicklerObjectTests, PickleBigmemPickleTests, ) test_support.run_doctest(pickle) + if os.path.exists(test_support.TESTFN): + os.unlink(test_support.TESTFN) if __name__ == "__main__": test_main()