Packing ExtType and some cleanup

7 files changed, 132 insertions, 170 deletions

diff --git a/msgpack/__init__.py b/msgpack/__init__.py
index 79107b6..a7b47b1 100644
--- a/msgpack/__init__.py
+++ b/msgpack/__init__.py
@@ -26,6 +26,7 @@ def pack(o, stream, **kwargs):
     packer = Packer(**kwargs)
     stream.write(packer.pack(o))
 
+
 def packb(o, **kwargs):
     """
     Pack object `o` and return packed bytes
@@ -40,4 +41,3 @@ loads = unpackb
 
 dump = pack
 dumps = packb
-
diff --git a/msgpack/_packer.pyx b/msgpack/_packer.pyx
index f033263..f63667c 100644
--- a/msgpack/_packer.pyx
+++ b/msgpack/_packer.pyx
@@ -8,6 +8,8 @@ from libc.limits cimport *
 from libc.stdint cimport int8_t
 
 from msgpack.exceptions import PackValueError
+from msgpack import ExtType
+
 
 cdef extern from "pack.h":
     struct msgpack_packer:
@@ -120,80 +122,87 @@ cdef class Packer(object):
         cdef int ret
         cdef dict d
         cdef size_t L
+        cdef int default_used = 0
 
         if nest_limit < 0:
             raise PackValueError("recursion limit exceeded.")
 
-        if o is None:
-            ret = msgpack_pack_nil(&self.pk)
-        elif isinstance(o, bool):
-            if o:
-                ret = msgpack_pack_true(&self.pk)
-            else:
-                ret = msgpack_pack_false(&self.pk)
-        elif PyLong_Check(o):
-            if o > 0:
-                ullval = o
-                ret = msgpack_pack_unsigned_long_long(&self.pk, ullval)
-            else:
-                llval = o
-                ret = msgpack_pack_long_long(&self.pk, llval)
-        elif PyInt_Check(o):
-            longval = o
-            ret = msgpack_pack_long(&self.pk, longval)
-        elif PyFloat_Check(o):
-            if self.use_float:
-               fval = o
-               ret = msgpack_pack_float(&self.pk, fval)
-            else:
-               dval = o
-               ret = msgpack_pack_double(&self.pk, dval)
-        elif PyBytes_Check(o):
-            rawval = o
-            L = len(o)
-            ret = msgpack_pack_bin(&self.pk, L)
-            if ret == 0:
+        while True:
+            if o is None:
+                ret = msgpack_pack_nil(&self.pk)
+            elif isinstance(o, bool):
+                if o:
+                    ret = msgpack_pack_true(&self.pk)
+                else:
+                    ret = msgpack_pack_false(&self.pk)
+            elif PyLong_Check(o):
+                if o > 0:
+                    ullval = o
+                    ret = msgpack_pack_unsigned_long_long(&self.pk, ullval)
+                else:
+                    llval = o
+                    ret = msgpack_pack_long_long(&self.pk, llval)
+            elif PyInt_Check(o):
+                longval = o
+                ret = msgpack_pack_long(&self.pk, longval)
+            elif PyFloat_Check(o):
+                if self.use_float:
+                   fval = o
+                   ret = msgpack_pack_float(&self.pk, fval)
+                else:
+                   dval = o
+                   ret = msgpack_pack_double(&self.pk, dval)
+            elif PyBytes_Check(o):
+                rawval = o
+                L = len(o)
+                ret = msgpack_pack_bin(&self.pk, L)
+                if ret == 0:
+                    ret = msgpack_pack_raw_body(&self.pk, rawval, L)
+            elif PyUnicode_Check(o):
+                if not self.encoding:
+                    raise TypeError("Can't encode unicode string: no encoding is specified")
+                o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
+                rawval = o
+                ret = msgpack_pack_raw(&self.pk, len(o))
+                if ret == 0:
+                    ret = msgpack_pack_raw_body(&self.pk, rawval, len(o))
+            elif PyDict_CheckExact(o):
+                d = <dict>o
+                ret = msgpack_pack_map(&self.pk, len(d))
+                if ret == 0:
+                    for k, v in d.iteritems():
+                        ret = self._pack(k, nest_limit-1)
+                        if ret != 0: break
+                        ret = self._pack(v, nest_limit-1)
+                        if ret != 0: break
+            elif PyDict_Check(o):
+                ret = msgpack_pack_map(&self.pk, len(o))
+                if ret == 0:
+                    for k, v in o.items():
+                        ret = self._pack(k, nest_limit-1)
+                        if ret != 0: break
+                        ret = self._pack(v, nest_limit-1)
+                        if ret != 0: break
+            elif isinstance(o, ExtType):
+                # This should be before Tuple because ExtType is namedtuple.
+                longval = o[0]
+                rawval = o[1]
+                L = len(o[1])
+                ret = msgpack_pack_ext(&self.pk, longval, L)
                 ret = msgpack_pack_raw_body(&self.pk, rawval, L)
-        elif PyUnicode_Check(o):
-            if not self.encoding:
-                raise TypeError("Can't encode unicode string: no encoding is specified")
-            o = PyUnicode_AsEncodedString(o, self.encoding, self.unicode_errors)
-            rawval = o
-            ret = msgpack_pack_raw(&self.pk, len(o))
-            if ret == 0:
-                ret = msgpack_pack_raw_body(&self.pk, rawval, len(o))
-        elif PyDict_CheckExact(o):
-            d = <dict>o
-            ret = msgpack_pack_map(&self.pk, len(d))
-            if ret == 0:
-                for k, v in d.iteritems():
-                    ret = self._pack(k, nest_limit-1)
-                    if ret != 0: break
-                    ret = self._pack(v, nest_limit-1)
-                    if ret != 0: break
-        elif PyDict_Check(o):
-            ret = msgpack_pack_map(&self.pk, len(o))
-            if ret == 0:
-                for k, v in o.items():
-                    ret = self._pack(k, nest_limit-1)
-                    if ret != 0: break
-                    ret = self._pack(v, nest_limit-1)
-                    if ret != 0: break
-        elif PyTuple_Check(o) or PyList_Check(o):
-            ret = msgpack_pack_array(&self.pk, len(o))
-            if ret == 0:
-                for v in o:
-                    ret = self._pack(v, nest_limit-1)
-                    if ret != 0: break
-        elif self.handle_unknown_type(o):
-            # it means that obj was succesfully packed, so we are done
-            return 0
-        elif self._default:
-            o = self._default(o)
-            ret = self._pack(o, nest_limit-1)
-        else:
-            raise TypeError("can't serialize %r" % (o,))
-        return ret
+            elif PyTuple_Check(o) or PyList_Check(o):
+                ret = msgpack_pack_array(&self.pk, len(o))
+                if ret == 0:
+                    for v in o:
+                        ret = self._pack(v, nest_limit-1)
+                        if ret != 0: break
+            elif not default_used and self._default:
+                o = self._default(o)
+                default_used = 1
+                continue
+            else:
+                raise TypeError("can't serialize %r" % (o,))
+            return ret
 
     cpdef pack(self, object obj):
         cdef int ret
@@ -207,9 +216,6 @@ cdef class Packer(object):
             self.pk.length = 0
             return buf
 
-    def handle_unknown_type(self, obj):
-        return None
-
     def pack_ext_type(self, typecode, data):
         msgpack_pack_ext(&self.pk, typecode, len(data))
         msgpack_pack_raw_body(&self.pk, data, len(data))
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
index b0e66db..d5aa46e 100644
--- a/msgpack/_unpacker.pyx
+++ b/msgpack/_unpacker.pyx
@@ -16,6 +16,7 @@ from msgpack.exceptions import (
         UnpackValueError,
         ExtraData,
         )
+from msgpack import ExtType
 
 
 cdef extern from "unpack.h":
@@ -24,7 +25,7 @@ cdef extern from "unpack.h":
         PyObject* object_hook
         bint has_pairs_hook # call object_hook with k-v pairs
         PyObject* list_hook
-        PyObject* ext_type_hook
+        PyObject* ext_hook
         char *encoding
         char *unicode_errors
 
@@ -43,8 +44,8 @@ cdef extern from "unpack.h":
     object unpack_data(unpack_context* ctx)
 
 cdef inline init_ctx(unpack_context *ctx,
-                     object object_hook, object object_pairs_hook, object list_hook,
-                     object ext_type_hook,
+                     object object_hook, object object_pairs_hook,
+                     object list_hook, object ext_hook,
                      bint use_list, char* encoding, char* unicode_errors):
     unpack_init(ctx)
     ctx.user.use_list = use_list
@@ -71,10 +72,10 @@ cdef inline init_ctx(unpack_context *ctx,
             raise TypeError("list_hook must be a callable.")
         ctx.user.list_hook = <PyObject*>list_hook
 
-    if ext_type_hook is not None:
-        if not PyCallable_Check(ext_type_hook):
-            raise TypeError("ext_type_hook must be a callable.")
-        ctx.user.ext_type_hook = <PyObject*>ext_type_hook
+    if ext_hook is not None:
+        if not PyCallable_Check(ext_hook):
+            raise TypeError("ext_hook must be a callable.")
+        ctx.user.ext_hook = <PyObject*>ext_hook
 
     ctx.user.encoding = encoding
     ctx.user.unicode_errors = unicode_errors
@@ -84,8 +85,7 @@ def default_read_extended_type(typecode, data):
 
 def unpackb(object packed, object object_hook=None, object list_hook=None,
             bint use_list=1, encoding=None, unicode_errors="strict",
-            object_pairs_hook=None,
-            ):
+            object_pairs_hook=None, ext_hook=ExtType):
     """
     Unpack packed_bytes to object. Returns an unpacked object.
 
@@ -114,8 +114,8 @@ def unpackb(object packed, object object_hook=None, object list_hook=None,
             unicode_errors = unicode_errors.encode('ascii')
         cerr = PyBytes_AsString(unicode_errors)
 
-    init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, default_read_extended_type,
-              use_list, cenc, cerr)
+    init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
+             use_list, cenc, cerr)
     ret = unpack_construct(&ctx, buf, buf_len, &off)
     if ret == 1:
         obj = unpack_data(&ctx)
@@ -220,7 +220,7 @@ cdef class Unpacker(object):
     def __init__(self, file_like=None, Py_ssize_t read_size=0, bint use_list=1,
                  object object_hook=None, object object_pairs_hook=None, object list_hook=None,
                  str encoding=None, str unicode_errors='strict', int max_buffer_size=0,
-                 ):
+                 object ext_hook=ExtType):
         cdef char *cenc=NULL, *cerr=NULL
 
         self.file_like = file_like
@@ -257,10 +257,8 @@ cdef class Unpacker(object):
                 self.unicode_errors = unicode_errors
             cerr = PyBytes_AsString(self.unicode_errors)
 
-        ext_type_hook = self.read_extended_type
-        Py_INCREF(ext_type_hook)
         init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook,
-                  ext_type_hook, use_list, cenc, cerr)
+                 ext_hook, use_list, cenc, cerr)
 
     def feed(self, object next_bytes):
         """Append `next_bytes` to internal buffer."""
@@ -370,24 +368,6 @@ cdef class Unpacker(object):
         """
         return self._unpack(unpack_construct, write_bytes)
 
-    def unpack_one(self, object write_bytes=None):
-        """
-        unpack one object
-
-        If write_bytes is not None, it will be called with parts of the raw
-        message as it is unpacked.
-
-        Raises `UnpackValueError` if there are no more bytes to unpack.
-        Raises ``ExtraData`` if there are still bytes left after the unpacking.
-        """
-        try:
-            result = self.unpack()
-        except OutOfData:
-            raise UnpackValueError("Data is not enough")
-        if self.buf_head < self.buf_tail:
-            raise ExtraData(result, self.buf[self.buf_head:])
-        return result
-
     def skip(self, object write_bytes=None):
         """
         read and ignore one object, returning None
@@ -415,9 +395,6 @@ cdef class Unpacker(object):
         """
         return self._unpack(read_map_header, write_bytes)
 
-    def read_extended_type(self, typecode, data):
-        return default_read_extended_type(typecode, data)
-
     def __iter__(self):
         return self
 
diff --git a/msgpack/pack_template.h b/msgpack/pack_template.h
index 0fe9514..8b91619 100644
--- a/msgpack/pack_template.h
+++ b/msgpack/pack_template.h
@@ -687,7 +687,7 @@ static inline int msgpack_pack_raw(msgpack_packer* x, size_t l)
 static inline int msgpack_pack_bin(msgpack_packer *x, size_t l)
 {
     if (!x->use_bin_type) {
-        return msgpack_pack_raw(x, l)
+        return msgpack_pack_raw(x, l);
     }
     if (l < 256) {
         unsigned char buf[2] = {0xc4, (unsigned char)l};
@@ -711,7 +711,6 @@ static inline int msgpack_pack_raw_body(msgpack_packer* x, const void* b, size_t
 /*
  * Ext
  */
-
 static inline int msgpack_pack_ext(msgpack_packer* x, int8_t typecode, size_t l)
 {
     if (l == 1) {
diff --git a/msgpack/unpack.h b/msgpack/unpack.h
index 327a524..3c09747 100644
--- a/msgpack/unpack.h
+++ b/msgpack/unpack.h
@@ -24,7 +24,7 @@ typedef struct unpack_user {
     PyObject *object_hook;
     bool has_pairs_hook;
     PyObject *list_hook;
-    PyObject *ext_type_hook;
+    PyObject *ext_hook;
     const char *encoding;
     const char *unicode_errors;
 } unpack_user;
@@ -241,12 +241,12 @@ static inline int unpack_callback_ext(unpack_user* u, const char* base, const ch
 {
     PyObject *py;
     int8_t typecode = (int8_t)*pos++;
-    if (!u->ext_type_hook) {
-        PyErr_SetString(PyExc_AssertionError, "u->ext_type_hook cannot be NULL");
+    if (!u->ext_hook) {
+        PyErr_SetString(PyExc_AssertionError, "u->ext_hook cannot be NULL");
         return -1;
     }
-    // lenght also includes the typecode, so the actual data is lenght-1
-    py = PyEval_CallFunction(u->ext_type_hook, "(is#)", typecode, pos, lenght-1);
+    // length also includes the typecode, so the actual data is lenght-1
+    py = PyEval_CallFunction(u->ext_hook, "(is#)", typecode, pos, lenght-1);
     if (!py)
         return -1;
     *o = py;
diff --git a/test/test_extension.py b/test/test_extension.py
index f2fa363..2f85ce3 100644
--- a/test/test_extension.py
+++ b/test/test_extension.py
@@ -1,5 +1,7 @@
+from __future__ import print_function
 import array
 import msgpack
+from msgpack import ExtType
 
 
 def test_pack_ext_type():
@@ -11,54 +13,45 @@ def test_pack_ext_type():
     assert p(b'AB')       == b'\xd5\x42AB'         # fixext 2
     assert p(b'ABCD')     == b'\xd6\x42ABCD'       # fixext 4
     assert p(b'ABCDEFGH') == b'\xd7\x42ABCDEFGH'   # fixext 8
-    assert p(b'A'*16)     == b'\xd8\x42' + 'A'*16  # fixext 16
+    assert p(b'A'*16)     == b'\xd8\x42' + b'A'*16 # fixext 16
     assert p(b'ABC')      == b'\xc7\x03\x42ABC'        # ext 8
     assert p(b'A'*0x0123)     == b'\xc8\x01\x23\x42' + b'A'*0x0123 # ext 16
     assert p(b'A'*0x00012345) == b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345 # ext 32
 
 
-def test_unpack_extended_type():
-    class MyUnpacker(msgpack.Unpacker):
-        def read_extended_type(self, typecode, data):
-            return (typecode, data)
+def test_unpack_ext_type():
+    def check(b, expected):
+        assert msgpack.unpackb(b) == expected
 
-    def u(s):
-        unpacker = MyUnpacker()
-        unpacker.feed(s)
-        return unpacker.unpack_one()
-
-    assert u('\xd4\x42A')         == (0x42, 'A')        # fixext 1
-    assert u('\xd5\x42AB')        == (0x42, 'AB')       # fixext 2
-    assert u('\xd6\x42ABCD')      == (0x42, 'ABCD')     # fixext 4
-    assert u('\xd7\x42ABCDEFGH')  == (0x42, 'ABCDEFGH') # fixext 8
-    assert u('\xd8\x42' + 'A'*16) == (0x42, 'A'*16)     # fixext 16
-    assert u('\xc7\x03\x42ABC')   == (0x42, 'ABC')      # ext 8
-    assert (u('\xc8\x01\x23\x42' + 'A'*0x0123) ==
-            (0x42, 'A'*0x0123))                         # ext 16
-    assert (u('\xc9\x00\x01\x23\x45\x42' + 'A'*0x00012345) ==
-            (0x42, 'A'*0x00012345))                     # ext 32
+    check(b'\xd4\x42A',         ExtType(0x42, b'A'))        # fixext 1
+    check(b'\xd5\x42AB',        ExtType(0x42, b'AB'))       # fixext 2
+    check(b'\xd6\x42ABCD',      ExtType(0x42, b'ABCD'))     # fixext 4
+    check(b'\xd7\x42ABCDEFGH',  ExtType(0x42, b'ABCDEFGH')) # fixext 8
+    check(b'\xd8\x42' + b'A'*16, ExtType(0x42, b'A'*16))    # fixext 16
+    check(b'\xc7\x03\x42ABC',   ExtType(0x42, b'ABC'))      # ext 8
+    check(b'\xc8\x01\x23\x42' + b'A'*0x0123,
+          ExtType(0x42, b'A'*0x0123))                        # ext 16
+    check(b'\xc9\x00\x01\x23\x45\x42' + b'A'*0x00012345,
+          ExtType(0x42, b'A'*0x00012345))                   # ext 32
 
 
 def test_extension_type():
-    class MyPacker(msgpack.Packer):
-        def handle_unknown_type(self, obj):
-            if isinstance(obj, array.array):
-                typecode = 123 # application specific typecode
-                data = obj.tostring()
-                self.pack_ext_type(typecode, data)
-                return True
-
-    class MyUnpacker(msgpack.Unpacker):
-        def read_extended_type(self, typecode, data):
-            assert typecode == 123
-            obj = array.array('d')
-            obj.fromstring(data)
-            return obj
-
-    obj = [42, 'hello', array.array('d', [1.1, 2.2, 3.3])]
-    packer = MyPacker()
-    unpacker = MyUnpacker(None)
-    s = packer.pack(obj)
-    unpacker.feed(s)
-    obj2 = unpacker.unpack_one()
+    def default(obj):
+        print('default called', obj)
+        if isinstance(obj, array.array):
+            typecode = 123 # application specific typecode
+            data = obj.tostring()
+            return ExtType(typecode, data)
+        raise TypeError("Unknwon type object %r" % (obj,))
+
+    def ext_hook(code, data):
+        print('ext_hook called', code, data)
+        assert code == 123
+        obj = array.array('d')
+        obj.fromstring(data)
+        return obj
+
+    obj = [42, b'hello', array.array('d', [1.1, 2.2, 3.3])]
+    s = msgpack.packb(obj, default=default)
+    obj2 = msgpack.unpackb(s, ext_hook=ext_hook)
     assert obj == obj2
diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py
index abc447a..af66b78 100644
--- a/test/test_sequnpack.py
+++ b/test/test_sequnpack.py
@@ -85,16 +85,3 @@ def test_readbytes():
     assert unpacker.read_bytes(3) == b'oob'
     assert unpacker.unpack() == ord(b'a')
     assert unpacker.unpack() == ord(b'r')
-
-def test_unpack_one():
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03abc')
-    assert unpacker.unpack_one() == 'abc'
-    #
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03abcd')
-    py.test.raises(ExtraData, "unpacker.unpack_one()")
-    #
-    unpacker = Unpacker()
-    unpacker.feed('\xda\x00\x03ab')
-    py.test.raises(UnpackValueError, "unpacker.unpack_one()")