Major speedup on packing.

2 files changed, 39 insertions, 50 deletions

diff --git a/msgpack/_msgpack.pyx b/msgpack/_msgpack.pyx
index 26222aa..9d766c5 100644
--- a/msgpack/_msgpack.pyx
+++ b/msgpack/_msgpack.pyx
@@ -18,6 +18,7 @@ cdef extern from "Python.h":
     int PyFloat_Check(object o)
     int PyString_Check(object o)
     int PyUnicode_Check(object o)
+    object PyBuffer_FromMemory(const_char_ptr b, Py_ssize_t len)
 
 cdef extern from "stdlib.h":
     void* malloc(size_t)
@@ -30,7 +31,9 @@ cdef extern from "string.h":
 
 cdef extern from "pack.h":
     struct msgpack_packer:
-        PyObject* writer
+        char* buf
+        size_t length
+        size_t buf_size
 
     int msgpack_pack_int(msgpack_packer* pk, int d)
     int msgpack_pack_nil(msgpack_packer* pk)
@@ -56,45 +59,16 @@ cdef class Packer(object):
     cdef object strm
     cdef object writer
 
-    def __init__(self, strm_, int size=4*1024):
-        self.strm = strm_
-        self.writer = strm_.write
-        self.pk.writer = <PyObject*>self.writer
+    def __init__(self, strm):
+        self.strm = strm
 
-    def flush(self):
-        """Flash local buffer and output stream if it has 'flush()' method."""
-        if hasattr(self.strm, 'flush'):
-            self.strm.flush()
-
-    def pack_list(self, len):
-        """Start packing sequential objects.
-
-        Example:
-
-            packer.pack_list(2)
-            packer.pack('foo')
-            packer.pack('bar')
-
-        This is same to:
-
-            packer.pack(['foo', 'bar'])
-        """
-        msgpack_pack_array(&self.pk, len)
+        cdef int buf_size = 1024*1024
+        self.pk.buf = <char*> malloc(buf_size);
+        self.pk.buf_size = buf_size
+        self.pk.length = 0
 
-    def pack_dict(self, len):
-        """Start packing key-value objects.
-
-        Example:
-
-            packer.pack_dict(1)
-            packer.pack('foo')
-            packer.pack('bar')
-
-        This is same to:
-
-            packer.pack({'foo': 'bar'})
-        """
-        msgpack_pack_map(&self.pk, len)
+    def __del__(self):
+        free(self.pk.buf);
 
     cdef __pack(self, object o):
         cdef long long llval
@@ -144,11 +118,19 @@ cdef class Packer(object):
             # TODO: Serialize with defalt() like simplejson.
             raise TypeError, "can't serialize %r" % (o,)
 
-    def pack(self, object obj, flush=True):
+    def pack(self, object obj, object flush=True):
         self.__pack(obj)
+        buf = PyBuffer_FromMemory(self.pk.buf, self.pk.length)
+        self.pk.length = 0
+        self.strm.write(buf)
         if flush:
             self.flush()
 
+    def flush(self):
+        """Flash local buffer and output stream if it has 'flush()' method."""
+        if hasattr(self.strm, 'flush'):
+            self.strm.flush()
+
     close = flush
 
 def pack(object o, object stream):
diff --git a/msgpack/pack.h b/msgpack/pack.h
index ea97601..d7e0867 100644
--- a/msgpack/pack.h
+++ b/msgpack/pack.h
@@ -25,7 +25,9 @@ extern "C" {
 #endif
 
 typedef struct msgpack_packer {
-    PyObject* writer;
+    char *buf;
+    size_t length;
+    size_t buf_size;
 } msgpack_packer;
 
 typedef struct Packer Packer;
@@ -64,16 +66,21 @@ static inline int msgpack_pack_raw_body(msgpack_packer* pk, const void* b, size_
 
 static inline int msgpack_pack_write(msgpack_packer* pk, const char *data, size_t l)
 {
-    PyObject *buf, *ret;
-
-    buf = PyBuffer_FromMemory((void*)data, l);
-    //buf = PyString_FromStringAndSize(data, l);
-    if (buf == NULL) return -1;
-
-    ret = PyObject_CallFunctionObjArgs(pk->writer, buf, NULL);
-    Py_DECREF(buf);
-    if (ret == NULL) return -1;
-    Py_DECREF(ret);
+    char* buf = pk->buf;
+    size_t bs = pk->buf_size;
+    size_t len = pk->length;
+
+    if (len + l > bs) {
+        bs = (len + l) * 2;
+        buf = realloc(pk->buf, bs);
+        if (!buf) return -1;
+    }
+    memcpy(buf + len, data, l);
+    len += l;
+
+    pk->buf = buf;
+    pk->buf_size = bs;
+    pk->length = len;
     return 0;
 }