Unpacker's ext_hook fixed + tests

Alexey Popravka · Alexey Popravka · commit d850e56dd0e2 · 2014-03-24T15:31:06.000+02:00
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
@@ -207,7 +207,7 @@ cdef class Unpacker(object):
     cdef object file_like_read
     cdef Py_ssize_t read_size
     # 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 size_t max_buffer_size
 
@@ -228,6 +228,7 @@ cdef class Unpacker(object):
         self.object_hook = object_hook
         self.object_pairs_hook = object_pairs_hook
         self.list_hook = list_hook
+        self.ext_hook = ext_hook
 
         self.file_like = file_like
         if file_like:
diff --git a/test/test_unpack.py b/test/test_unpack.py
@@ -1,6 +1,6 @@
 from io import BytesIO
 import sys
-from msgpack import Unpacker, packb, OutOfData
+from msgpack import Unpacker, packb, OutOfData, ExtType
 from pytest import raises, mark
 
 
@@ -42,6 +42,29 @@ def hook(x):
     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__':
     test_unpack_array_header_from_file()
     test_unpacker_hook_refcnt()
+    test_unpacker_ext_hook()
