test_cpickle.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257
  1. import cPickle
  2. import cStringIO
  3. import io
  4. import functools
  5. import unittest
  6. from test.pickletester import (AbstractUnpickleTests,
  7. AbstractPickleTests,
  8. AbstractPickleModuleTests,
  9. AbstractPicklerUnpicklerObjectTests,
  10. BigmemPickleTests)
  11. from test import test_support
  12. class cStringIOMixin:
  13. output = input = cStringIO.StringIO
  14. def close(self, f):
  15. pass
  16. class BytesIOMixin:
  17. output = input = io.BytesIO
  18. def close(self, f):
  19. pass
  20. class FileIOMixin:
  21. def output(self):
  22. return open(test_support.TESTFN, 'wb+')
  23. def input(self, data):
  24. f = open(test_support.TESTFN, 'wb+')
  25. try:
  26. f.write(data)
  27. f.seek(0)
  28. return f
  29. except:
  30. f.close()
  31. raise
  32. def close(self, f):
  33. f.close()
  34. test_support.unlink(test_support.TESTFN)
  35. class cPickleTests(AbstractUnpickleTests, AbstractPickleTests,
  36. AbstractPickleModuleTests):
  37. def setUp(self):
  38. self.dumps = cPickle.dumps
  39. self.loads = cPickle.loads
  40. error = cPickle.BadPickleGet
  41. module = cPickle
  42. bad_stack_errors = (cPickle.UnpicklingError,)
  43. bad_mark_errors = (EOFError,)
  44. truncated_errors = (cPickle.UnpicklingError, EOFError,
  45. AttributeError, ValueError)
  46. class cPickleUnpicklerTests(AbstractUnpickleTests):
  47. def loads(self, buf):
  48. f = self.input(buf)
  49. try:
  50. p = cPickle.Unpickler(f)
  51. return p.load()
  52. finally:
  53. self.close(f)
  54. error = cPickle.BadPickleGet
  55. bad_stack_errors = (cPickle.UnpicklingError,)
  56. bad_mark_errors = (EOFError,)
  57. truncated_errors = (cPickle.UnpicklingError, EOFError,
  58. AttributeError, ValueError)
  59. class cStringIOCUnpicklerTests(cStringIOMixin, cPickleUnpicklerTests):
  60. pass
  61. class BytesIOCUnpicklerTests(BytesIOMixin, cPickleUnpicklerTests):
  62. pass
  63. class FileIOCUnpicklerTests(FileIOMixin, cPickleUnpicklerTests):
  64. pass
  65. class cPicklePicklerTests(AbstractPickleTests):
  66. def dumps(self, arg, proto=0):
  67. f = self.output()
  68. try:
  69. p = cPickle.Pickler(f, proto)
  70. p.dump(arg)
  71. f.seek(0)
  72. return f.read()
  73. finally:
  74. self.close(f)
  75. def loads(self, buf):
  76. f = self.input(buf)
  77. try:
  78. p = cPickle.Unpickler(f)
  79. return p.load()
  80. finally:
  81. self.close(f)
  82. class cStringIOCPicklerTests(cStringIOMixin, cPicklePicklerTests):
  83. pass
  84. class BytesIOCPicklerTests(BytesIOMixin, cPicklePicklerTests):
  85. pass
  86. class FileIOCPicklerTests(FileIOMixin, cPicklePicklerTests):
  87. pass
  88. class cPickleListPicklerTests(AbstractPickleTests):
  89. def dumps(self, arg, proto=0):
  90. p = cPickle.Pickler(proto)
  91. p.dump(arg)
  92. return p.getvalue()
  93. def loads(self, *args):
  94. f = self.input(args[0])
  95. try:
  96. p = cPickle.Unpickler(f)
  97. return p.load()
  98. finally:
  99. self.close(f)
  100. error = cPickle.BadPickleGet
  101. class cStringIOCPicklerListTests(cStringIOMixin, cPickleListPicklerTests):
  102. pass
  103. class BytesIOCPicklerListTests(BytesIOMixin, cPickleListPicklerTests):
  104. pass
  105. class FileIOCPicklerListTests(FileIOMixin, cPickleListPicklerTests):
  106. pass
  107. class cPickleFastPicklerTests(AbstractPickleTests):
  108. def dumps(self, arg, proto=0):
  109. f = self.output()
  110. try:
  111. p = cPickle.Pickler(f, proto)
  112. p.fast = 1
  113. p.dump(arg)
  114. f.seek(0)
  115. return f.read()
  116. finally:
  117. self.close(f)
  118. def loads(self, *args):
  119. f = self.input(args[0])
  120. try:
  121. p = cPickle.Unpickler(f)
  122. return p.load()
  123. finally:
  124. self.close(f)
  125. def test_nonrecursive_deep(self):
  126. # If it's not cyclic, it should pickle OK even if the nesting
  127. # depth exceeds PY_CPICKLE_FAST_LIMIT. That happens to be
  128. # 50 today. Jack Jansen reported stack overflow on Mac OS 9
  129. # at 64.
  130. a = []
  131. for i in range(60):
  132. a = [a]
  133. b = self.loads(self.dumps(a))
  134. self.assertEqual(a, b)
  135. for name in dir(AbstractPickleTests):
  136. if name.startswith('test_recursive_'):
  137. func = getattr(AbstractPickleTests, name)
  138. if '_subclass' in name and '_and_inst' not in name:
  139. assert_args = RuntimeError, 'maximum recursion depth exceeded'
  140. else:
  141. assert_args = ValueError, "can't pickle cyclic objects"
  142. def wrapper(self, func=func, assert_args=assert_args):
  143. with self.assertRaisesRegexp(*assert_args):
  144. func(self)
  145. functools.update_wrapper(wrapper, func)
  146. setattr(cPickleFastPicklerTests, name, wrapper)
  147. class cStringIOCPicklerFastTests(cStringIOMixin, cPickleFastPicklerTests):
  148. pass
  149. class BytesIOCPicklerFastTests(BytesIOMixin, cPickleFastPicklerTests):
  150. pass
  151. class FileIOCPicklerFastTests(FileIOMixin, cPickleFastPicklerTests):
  152. pass
  153. class cPicklePicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):
  154. pickler_class = cPickle.Pickler
  155. unpickler_class = cPickle.Unpickler
  156. class cPickleBigmemPickleTests(BigmemPickleTests):
  157. def dumps(self, arg, proto=0, fast=0):
  158. # Ignore fast
  159. return cPickle.dumps(arg, proto)
  160. def loads(self, buf):
  161. # Ignore fast
  162. return cPickle.loads(buf)
  163. class Node(object):
  164. pass
  165. class cPickleDeepRecursive(unittest.TestCase):
  166. def test_issue2702(self):
  167. # This should raise a RecursionLimit but in some
  168. # platforms (FreeBSD, win32) sometimes raises KeyError instead,
  169. # or just silently terminates the interpreter (=crashes).
  170. nodes = [Node() for i in range(500)]
  171. for n in nodes:
  172. n.connections = list(nodes)
  173. n.connections.remove(n)
  174. self.assertRaises((AttributeError, RuntimeError), cPickle.dumps, n)
  175. def test_issue3179(self):
  176. # Safe test, because I broke this case when fixing the
  177. # behaviour for the previous test.
  178. res=[]
  179. for x in range(1,2000):
  180. res.append(dict(doc=x, similar=[]))
  181. cPickle.dumps(res)
  182. def test_main():
  183. test_support.run_unittest(
  184. cPickleTests,
  185. cStringIOCUnpicklerTests,
  186. BytesIOCUnpicklerTests,
  187. FileIOCUnpicklerTests,
  188. cStringIOCPicklerTests,
  189. BytesIOCPicklerTests,
  190. FileIOCPicklerTests,
  191. cStringIOCPicklerListTests,
  192. BytesIOCPicklerListTests,
  193. FileIOCPicklerListTests,
  194. cStringIOCPicklerFastTests,
  195. BytesIOCPicklerFastTests,
  196. FileIOCPicklerFastTests,
  197. cPickleDeepRecursive,
  198. cPicklePicklerUnpicklerObjectTests,
  199. cPickleBigmemPickleTests,
  200. )
  201. if __name__ == "__main__":
  202. test_main()