diff --git a/python/msgpack/_msgpack.pyx b/python/msgpack/_msgpack.pyx index 26222aae..9d766c5a 100644 --- a/python/msgpack/_msgpack.pyx +++ b/python/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 = 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() + cdef int buf_size = 1024*1024 + self.pk.buf = malloc(buf_size); + self.pk.buf_size = buf_size + self.pk.length = 0 - 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) - - 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/python/msgpack/pack.h b/python/msgpack/pack.h index ea97601d..d7e08677 100644 --- a/python/msgpack/pack.h +++ b/python/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; + char* buf = pk->buf; + size_t bs = pk->buf_size; + size_t len = pk->length; - buf = PyBuffer_FromMemory((void*)data, l); - //buf = PyString_FromStringAndSize(data, l); - if (buf == NULL) return -1; + if (len + l > bs) { + bs = (len + l) * 2; + buf = realloc(pk->buf, bs); + if (!buf) return -1; + } + memcpy(buf + len, data, l); + len += l; - ret = PyObject_CallFunctionObjArgs(pk->writer, buf, NULL); - Py_DECREF(buf); - if (ret == NULL) return -1; - Py_DECREF(ret); + pk->buf = buf; + pk->buf_size = bs; + pk->length = len; return 0; }