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
  10 class PickleTests(AbstractPickleTests, AbstractPickleModuleTests):
  11
  12     def dumps(self, arg, proto=0, fast=0):
  13         # Ignore fast
  14         return pickle.dumps(arg, proto)
  15
  16     def loads(self, buf):
  17         # Ignore fast
  18         return pickle.loads(buf)
  19
  20     module = pickle
  21     error = KeyError
  22
  23 class PicklerTests(AbstractPickleTests):
  24
  25     error = KeyError
  26
  27     def dumps(self, arg, proto=0, fast=0):
  28         f = StringIO()
  29         p = pickle.Pickler(f, proto)
  30         if fast:
  31             p.fast = fast
  32         p.dump(arg)
  33         f.seek(0)
  34         return f.read()
  35
  36     def loads(self, buf):
  37         f = StringIO(buf)
  38         u = pickle.Unpickler(f)
  39         return u.load()
  40
  41 class PersPicklerTests(AbstractPersistentPicklerTests):
  42
  43     def dumps(self, arg, proto=0, fast=0):
  44         class PersPickler(pickle.Pickler):
  45             def persistent_id(subself, obj):
  46                 return self.persistent_id(obj)
  47         f = StringIO()
  48         p = PersPickler(f, proto)
  49         if fast:
  50             p.fast = fast
  51         p.dump(arg)
  52         f.seek(0)
  53         return f.read()
  54
  55     def loads(self, buf):
  56         class PersUnpickler(pickle.Unpickler):
  57             def persistent_load(subself, obj):
  58                 return self.persistent_load(obj)
  59         f = StringIO(buf)
  60         u = PersUnpickler(f)
  61         return u.load()
  62
  63 def test_main():
  64     test_support.run_unittest(
  65         PickleTests,
  66         PicklerTests,
  67         PersPicklerTests
  68     )
  69     test_support.run_doctest(pickle)
  70
  71 if __name__ == "__main__":
  72     test_main()