Merge pull request #93 from popravich/unpacker_ext_hook_fix

Unpacker's ext_hook fixed + tests
This commit is contained in:
INADA Naoki
2014-03-25 12:19:51 +09:00
2 changed files with 27 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,30 @@ 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(MyUnpacker, self).__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()