test_obj.py 2.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. # coding: utf-8
  2. import pytest
  3. from pandas.io.msgpack import packb, unpackb
  4. class DecodeError(Exception):
  5. pass
  6. class TestObj(object):
  7. def _arr_to_str(self, arr):
  8. return ''.join(str(c) for c in arr)
  9. def bad_complex_decoder(self, o):
  10. raise DecodeError("Ooops!")
  11. def _decode_complex(self, obj):
  12. if b'__complex__' in obj:
  13. return complex(obj[b'real'], obj[b'imag'])
  14. return obj
  15. def _encode_complex(self, obj):
  16. if isinstance(obj, complex):
  17. return {b'__complex__': True, b'real': 1, b'imag': 2}
  18. return obj
  19. def test_encode_hook(self):
  20. packed = packb([3, 1 + 2j], default=self._encode_complex)
  21. unpacked = unpackb(packed, use_list=1)
  22. assert unpacked[1] == {b'__complex__': True, b'real': 1, b'imag': 2}
  23. def test_decode_hook(self):
  24. packed = packb([3, {b'__complex__': True, b'real': 1, b'imag': 2}])
  25. unpacked = unpackb(packed, object_hook=self._decode_complex,
  26. use_list=1)
  27. assert unpacked[1] == 1 + 2j
  28. def test_decode_pairs_hook(self):
  29. packed = packb([3, {1: 2, 3: 4}])
  30. prod_sum = 1 * 2 + 3 * 4
  31. unpacked = unpackb(
  32. packed, object_pairs_hook=lambda l: sum(k * v for k, v in l),
  33. use_list=1)
  34. assert unpacked[1] == prod_sum
  35. def test_only_one_obj_hook(self):
  36. msg = "object_pairs_hook and object_hook are mutually exclusive"
  37. with pytest.raises(TypeError, match=msg):
  38. unpackb(b'', object_hook=lambda x: x,
  39. object_pairs_hook=lambda x: x)
  40. def test_bad_hook(self):
  41. msg = r"can't serialize \(1\+2j\)"
  42. with pytest.raises(TypeError, match=msg):
  43. packed = packb([3, 1 + 2j], default=lambda o: o)
  44. unpacked = unpackb(packed, use_list=1) # noqa
  45. def test_array_hook(self):
  46. packed = packb([1, 2, 3])
  47. unpacked = unpackb(packed, list_hook=self._arr_to_str, use_list=1)
  48. assert unpacked == '123'
  49. def test_an_exception_in_objecthook1(self):
  50. with pytest.raises(DecodeError, match='Ooops!'):
  51. packed = packb({1: {'__complex__': True, 'real': 1, 'imag': 2}})
  52. unpackb(packed, object_hook=self.bad_complex_decoder)
  53. def test_an_exception_in_objecthook2(self):
  54. with pytest.raises(DecodeError, match='Ooops!'):
  55. packed = packb({1: [{'__complex__': True, 'real': 1, 'imag': 2}]})
  56. unpackb(packed, list_hook=self.bad_complex_decoder, use_list=1)