test_pickle.py 3.3 KB

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