diff options
| -rw-r--r-- | msgpack/_unpacker.pyx | 7 | ||||
| -rw-r--r-- | test/test_unpack.py | 47 | ||||
| -rw-r--r-- | test/test_unpack_file.py | 19 |
3 files changed, 53 insertions, 20 deletions
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index 16aca5c..0df6ab3 100644 --- a/msgpack/_unpacker.pyx +++ b/msgpack/_unpacker.pyx @@ -206,7 +206,8 @@ cdef class Unpacker(object): cdef object file_like cdef object file_like_read cdef Py_ssize_t read_size - cdef object object_hook + # To maintain refcnt. + cdef object object_hook, object_pairs_hook, list_hook cdef object encoding, unicode_errors cdef size_t max_buffer_size @@ -224,6 +225,10 @@ cdef class Unpacker(object): cdef char *cenc=NULL, cdef char *cerr=NULL + self.object_hook = object_hook + self.object_pairs_hook = object_pairs_hook + self.list_hook = list_hook + self.file_like = file_like if file_like: self.file_like_read = file_like.read diff --git a/test/test_unpack.py b/test/test_unpack.py new file mode 100644 index 0000000..544cebf --- /dev/null +++ b/test/test_unpack.py @@ -0,0 +1,47 @@ +from io import BytesIO +import sys +from msgpack import Unpacker, packb, OutOfData +from pytest import raises, mark + + +def test_unpack_array_header_from_file(): + f = BytesIO(packb([1,2,3,4])) + unpacker = Unpacker(f) + assert unpacker.read_array_header() == 4 + assert unpacker.unpack() == 1 + assert unpacker.unpack() == 2 + assert unpacker.unpack() == 3 + assert unpacker.unpack() == 4 + with raises(OutOfData): + unpacker.unpack() + + +@mark.skipif(not hasattr(sys, 'getrefcount'), + reason='sys.getrefcount() is needed to pass this test') +def test_unpacker_hook_refcnt(): + result = [] + + def hook(x): + result.append(x) + return x + + basecnt = sys.getrefcount(hook) + + up = Unpacker(object_hook=hook, list_hook=hook) + + assert sys.getrefcount(hook) >= basecnt + 2 + + up.feed(packb([{}])) + up.feed(packb([{}])) + assert up.unpack() == [{}] + assert up.unpack() == [{}] + assert result == [{}, [{}], {}, [{}]] + + del up + + assert sys.getrefcount(hook) == basecnt + + +if __name__ == '__main__': + test_unpack_array_header_from_file() + test_unpacker_hook_refcnt() diff --git a/test/test_unpack_file.py b/test/test_unpack_file.py deleted file mode 100644 index 1563008..0000000 --- a/test/test_unpack_file.py +++ /dev/null @@ -1,19 +0,0 @@ -from io import BytesIO -from msgpack import Unpacker, packb, OutOfData -from pytest import raises - - -def test_unpack_array_header_from_file(): - f = BytesIO(packb([1,2,3,4])) - unpacker = Unpacker(f) - assert unpacker.read_array_header() == 4 - assert unpacker.unpack() == 1 - assert unpacker.unpack() == 2 - assert unpacker.unpack() == 3 - assert unpacker.unpack() == 4 - with raises(OutOfData): - unpacker.unpack() - - -if __name__ == '__main__': - test_unpack_array_header_from_file() |