Unpacker's ext_hook fixed + tests

This commit is contained in:
Alexey Popravka 2014-03-24 15:31:06 +02:00
parent e9de6b7f39
commit d850e56dd0
2 changed files with 26 additions and 2 deletions

View file

@ -207,7 +207,7 @@ cdef class Unpacker(object):
cdef object file_like_read cdef object file_like_read
cdef Py_ssize_t read_size cdef Py_ssize_t read_size
# To maintain refcnt. # To maintain refcnt.
cdef object object_hook, object_pairs_hook, list_hook cdef object object_hook, object_pairs_hook, list_hook, ext_hook
cdef object encoding, unicode_errors cdef object encoding, unicode_errors
cdef size_t max_buffer_size cdef size_t max_buffer_size
@ -228,6 +228,7 @@ cdef class Unpacker(object):
self.object_hook = object_hook self.object_hook = object_hook
self.object_pairs_hook = object_pairs_hook self.object_pairs_hook = object_pairs_hook
self.list_hook = list_hook self.list_hook = list_hook
self.ext_hook = ext_hook
self.file_like = file_like self.file_like = file_like
if file_like: if file_like:

View file

@ -1,6 +1,6 @@
from io import BytesIO from io import BytesIO
import sys import sys
from msgpack import Unpacker, packb, OutOfData from msgpack import Unpacker, packb, OutOfData, ExtType
from pytest import raises, mark from pytest import raises, mark
@ -42,6 +42,29 @@ def test_unpacker_hook_refcnt():
assert sys.getrefcount(hook) == basecnt assert sys.getrefcount(hook) == basecnt
def test_unpacker_ext_hook():
class MyUnpacker(Unpacker):
def __init__(self):
super().__init__(ext_hook=self._hook, encoding='utf-8')
def _hook(self, code, data):
if code == 1:
return int(data)
else:
return ExtType(code, data)
unpacker = MyUnpacker()
unpacker.feed(packb({'a': 1}, encoding='utf-8'))
assert unpacker.unpack() == {'a': 1}
unpacker.feed(packb({'a': ExtType(1, b'123')}, encoding='utf-8'))
assert unpacker.unpack() == {'a': 123}
unpacker.feed(packb({'a': ExtType(2, b'321')}, encoding='utf-8'))
assert unpacker.unpack() == {'a': ExtType(2, b'321')}
if __name__ == '__main__': if __name__ == '__main__':
test_unpack_array_header_from_file() test_unpack_array_header_from_file()
test_unpacker_hook_refcnt() test_unpacker_hook_refcnt()
test_unpacker_ext_hook()