Lib/test/test_pickle.py

   1 import pickle
   2 import unittest
   3 from cStringIO import StringIO
   4
   5 from test import test_support
   6
   7 from test.pickletester import AbstractPickleTests
   8 from test.pickletester import AbstractPickleModuleTests
   9 from test.pickletester import AbstractPersistentPicklerTests
  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 def test_main():
  65     loader = unittest.TestLoader()
  66     suite = unittest.TestSuite()
  67     suite.addTest(loader.loadTestsFromTestCase(PickleTests))
  68     suite.addTest(loader.loadTestsFromTestCase(PicklerTests))
  69     suite.addTest(loader.loadTestsFromTestCase(PersPicklerTests))
  70     test_support.run_suite(suite)
  71     test_support.run_doctest(pickle)
  72
  73 if __name__ == "__main__":
  74     test_main()