Lib/test/test_pickle.py

   1 import pickle
   2 from cStringIO import StringIO
   3
   4 from test import test_support
   5
   6 from test.pickletester import AbstractPickleTests
   7 from test.pickletester import AbstractPickleModuleTests
   8 from test.pickletester import AbstractPersistentPicklerTests
   9 from test.pickletester import AbstractPicklerUnpicklerObjectTests
  10
  11 class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
  12
  13     def dumps(self, arg, proto=0, fast=0):
  14         # Ignore fast
  15         return pickle.dumps(arg, proto)
  16
  17     def loads(self, buf):
  18         # Ignore fast
  19         return pickle.loads(buf)
  20
  21     module = pickle
  22     error = KeyError
  23
  24 class PicklerTests(AbstractPickleTests):
  25
  26     error = KeyError
  27
  28     def dumps(self, arg, proto=0, fast=0):
  29         f = StringIO()
  30         p = pickle.Pickler(f, proto)
  31         if fast:
  32             p.fast = fast
  33         p.dump(arg)
  34         f.seek(0)
  35         return f.read()
  36
  37     def loads(self, buf):
  38         f = StringIO(buf)
  39         u = pickle.Unpickler(f)
  40         return u.load()
  41
  42 class PersPicklerTests(AbstractPersistentPicklerTests):
  43
  44     def dumps(self, arg, proto=0, fast=0):
  45         class PersPickler(pickle.Pickler):
  46             def persistent_id(subself, obj):
  47                 return self.persistent_id(obj)
  48         f = StringIO()
  49         p = PersPickler(f, proto)
  50         if fast:
  51             p.fast = fast
  52         p.dump(arg)
  53         f.seek(0)
  54         return f.read()
  55
  56     def loads(self, buf):
  57         class PersUnpickler(pickle.Unpickler):
  58             def persistent_load(subself, obj):
  59                 return self.persistent_load(obj)
  60         f = StringIO(buf)
  61         u = PersUnpickler(f)
  62         return u.load()
  63
  64 class PicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
  65
  66     pickler_class = pickle.Pickler
  67     unpickler_class = pickle.Unpickler
  68
  69
  70 def test_main():
  71     test_support.run_unittest(
  72         PickleTests,
  73         PicklerTests,
  74         PersPicklerTests,
  75         PicklerUnpicklerObjectTests,
  76     )
  77     test_support.run_doctest(pickle)
  78
  79 if __name__ == "__main__":
  80     test_main()