test_picklebuffer.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. """Unit tests for the PickleBuffer object.
  2. Pickling tests themselves are in pickletester.py.
  3. """
  4. import gc
  5. from pickle import PickleBuffer
  6. import weakref
  7. import unittest
  8. from test.support import import_helper
  9. class B(bytes):
  10. pass
  11. class PickleBufferTest(unittest.TestCase):
  12. def check_memoryview(self, pb, equiv):
  13. with memoryview(pb) as m:
  14. with memoryview(equiv) as expected:
  15. self.assertEqual(m.nbytes, expected.nbytes)
  16. self.assertEqual(m.readonly, expected.readonly)
  17. self.assertEqual(m.itemsize, expected.itemsize)
  18. self.assertEqual(m.shape, expected.shape)
  19. self.assertEqual(m.strides, expected.strides)
  20. self.assertEqual(m.c_contiguous, expected.c_contiguous)
  21. self.assertEqual(m.f_contiguous, expected.f_contiguous)
  22. self.assertEqual(m.format, expected.format)
  23. self.assertEqual(m.tobytes(), expected.tobytes())
  24. def test_constructor_failure(self):
  25. with self.assertRaises(TypeError):
  26. PickleBuffer()
  27. with self.assertRaises(TypeError):
  28. PickleBuffer("foo")
  29. # Released memoryview fails taking a buffer
  30. m = memoryview(b"foo")
  31. m.release()
  32. with self.assertRaises(ValueError):
  33. PickleBuffer(m)
  34. def test_basics(self):
  35. pb = PickleBuffer(b"foo")
  36. self.assertEqual(b"foo", bytes(pb))
  37. with memoryview(pb) as m:
  38. self.assertTrue(m.readonly)
  39. pb = PickleBuffer(bytearray(b"foo"))
  40. self.assertEqual(b"foo", bytes(pb))
  41. with memoryview(pb) as m:
  42. self.assertFalse(m.readonly)
  43. m[0] = 48
  44. self.assertEqual(b"0oo", bytes(pb))
  45. def test_release(self):
  46. pb = PickleBuffer(b"foo")
  47. pb.release()
  48. with self.assertRaises(ValueError) as raises:
  49. memoryview(pb)
  50. self.assertIn("operation forbidden on released PickleBuffer object",
  51. str(raises.exception))
  52. # Idempotency
  53. pb.release()
  54. def test_cycle(self):
  55. b = B(b"foo")
  56. pb = PickleBuffer(b)
  57. b.cycle = pb
  58. wpb = weakref.ref(pb)
  59. del b, pb
  60. gc.collect()
  61. self.assertIsNone(wpb())
  62. def test_ndarray_2d(self):
  63. # C-contiguous
  64. ndarray = import_helper.import_module("_testbuffer").ndarray
  65. arr = ndarray(list(range(12)), shape=(4, 3), format='<i')
  66. self.assertTrue(arr.c_contiguous)
  67. self.assertFalse(arr.f_contiguous)
  68. pb = PickleBuffer(arr)
  69. self.check_memoryview(pb, arr)
  70. # Non-contiguous
  71. arr = arr[::2]
  72. self.assertFalse(arr.c_contiguous)
  73. self.assertFalse(arr.f_contiguous)
  74. pb = PickleBuffer(arr)
  75. self.check_memoryview(pb, arr)
  76. # F-contiguous
  77. arr = ndarray(list(range(12)), shape=(3, 4), strides=(4, 12), format='<i')
  78. self.assertTrue(arr.f_contiguous)
  79. self.assertFalse(arr.c_contiguous)
  80. pb = PickleBuffer(arr)
  81. self.check_memoryview(pb, arr)
  82. # Tests for PickleBuffer.raw()
  83. def check_raw(self, obj, equiv):
  84. pb = PickleBuffer(obj)
  85. with pb.raw() as m:
  86. self.assertIsInstance(m, memoryview)
  87. self.check_memoryview(m, equiv)
  88. def test_raw(self):
  89. for obj in (b"foo", bytearray(b"foo")):
  90. with self.subTest(obj=obj):
  91. self.check_raw(obj, obj)
  92. def test_raw_ndarray(self):
  93. # 1-D, contiguous
  94. ndarray = import_helper.import_module("_testbuffer").ndarray
  95. arr = ndarray(list(range(3)), shape=(3,), format='<h')
  96. equiv = b"\x00\x00\x01\x00\x02\x00"
  97. self.check_raw(arr, equiv)
  98. # 2-D, C-contiguous
  99. arr = ndarray(list(range(6)), shape=(2, 3), format='<h')
  100. equiv = b"\x00\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00"
  101. self.check_raw(arr, equiv)
  102. # 2-D, F-contiguous
  103. arr = ndarray(list(range(6)), shape=(2, 3), strides=(2, 4),
  104. format='<h')
  105. # Note this is different from arr.tobytes()
  106. equiv = b"\x00\x00\x01\x00\x02\x00\x03\x00\x04\x00\x05\x00"
  107. self.check_raw(arr, equiv)
  108. # 0-D
  109. arr = ndarray(456, shape=(), format='<i')
  110. equiv = b'\xc8\x01\x00\x00'
  111. self.check_raw(arr, equiv)
  112. def check_raw_non_contiguous(self, obj):
  113. pb = PickleBuffer(obj)
  114. with self.assertRaisesRegex(BufferError, "non-contiguous"):
  115. pb.raw()
  116. def test_raw_non_contiguous(self):
  117. # 1-D
  118. ndarray = import_helper.import_module("_testbuffer").ndarray
  119. arr = ndarray(list(range(6)), shape=(6,), format='<i')[::2]
  120. self.check_raw_non_contiguous(arr)
  121. # 2-D
  122. arr = ndarray(list(range(12)), shape=(4, 3), format='<i')[::2]
  123. self.check_raw_non_contiguous(arr)
  124. def test_raw_released(self):
  125. pb = PickleBuffer(b"foo")
  126. pb.release()
  127. with self.assertRaises(ValueError) as raises:
  128. pb.raw()
  129. if __name__ == "__main__":
  130. unittest.main()