Packing ExtType and some cleanup

msgpack · methane · Oct 20, 2013 · Oct 15, 2013 · Oct 18, 2013 · Oct 18, 2013
commit 96bcd76f49afd00f5b7def1ff7cfd002a7fa477d
diff --git 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 @@ def packb(o, **kwargs):
 
 dump = pack
 dumps = packb
-
diff --git 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
@@ -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
@@ -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
@@ -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;