cpython/Lib/test/test_pickle.py
2011-02-24 19:40:09 +00:00

117 lines
3.1 KiB
Python

import pickle
import io
from test import support
from test.pickletester import AbstractPickleTests
from test.pickletester import AbstractPickleModuleTests
from test.pickletester import AbstractPersistentPicklerTests
from test.pickletester import AbstractPicklerUnpicklerObjectTests
try:
import _pickle
has_c_implementation = True
except ImportError:
has_c_implementation = False
class PickleTests(AbstractPickleModuleTests):
pass
class PyPicklerTests(AbstractPickleTests):
pickler = pickle._Pickler
unpickler = pickle._Unpickler
def dumps(self, arg, proto=None):
f = io.BytesIO()
p = self.pickler(f, proto)
p.dump(arg)
f.seek(0)
return bytes(f.read())
def loads(self, buf, **kwds):
f = io.BytesIO(buf)
u = self.unpickler(f, **kwds)
return u.load()
class InMemoryPickleTests(AbstractPickleTests):
pickler = pickle._Pickler
unpickler = pickle._Unpickler
def dumps(self, arg, proto=None):
return pickle.dumps(arg, proto)
def loads(self, buf, **kwds):
return pickle.loads(buf, **kwds)
class PyPersPicklerTests(AbstractPersistentPicklerTests):
pickler = pickle._Pickler
unpickler = pickle._Unpickler
def dumps(self, arg, proto=None):
class PersPickler(self.pickler):
def persistent_id(subself, obj):
return self.persistent_id(obj)
f = io.BytesIO()
p = PersPickler(f, proto)
p.dump(arg)
f.seek(0)
return f.read()
def loads(self, buf, **kwds):
class PersUnpickler(self.unpickler):
def persistent_load(subself, obj):
return self.persistent_load(obj)
f = io.BytesIO(buf)
u = PersUnpickler(f, **kwds)
return u.load()
class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
pickler_class = pickle._Pickler
unpickler_class = pickle._Unpickler
if has_c_implementation:
class CPicklerTests(PyPicklerTests):
pickler = _pickle.Pickler
unpickler = _pickle.Unpickler
class CPersPicklerTests(PyPersPicklerTests):
pickler = _pickle.Pickler
unpickler = _pickle.Unpickler
class CDumpPickle_LoadPickle(PyPicklerTests):
pickler = _pickle.Pickler
unpickler = pickle._Unpickler
class DumpPickle_CLoadPickle(PyPicklerTests):
pickler = pickle._Pickler
unpickler = _pickle.Unpickler
class CPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
pickler_class = _pickle.Pickler
unpickler_class = _pickle.Unpickler
def test_main():
tests = [PickleTests, PyPicklerTests, PyPersPicklerTests]
if has_c_implementation:
tests.extend([CPicklerTests, CPersPicklerTests,
CDumpPickle_LoadPickle, DumpPickle_CLoadPickle,
PyPicklerUnpicklerObjectTests,
CPicklerUnpicklerObjectTests,
InMemoryPickleTests])
support.run_unittest(*tests)
support.run_doctest(pickle)
if __name__ == "__main__":
test_main()