python: Support old buffer protocol when unpack. (experimental)

This commit is contained in:
INADA Naoki
2010-11-03 03:11:00 +09:00
parent b1df5d3ad7
commit 4688252bd4
3 changed files with 65 additions and 113 deletions

View File

@@ -1,37 +1,14 @@
# coding: utf-8 # coding: utf-8
from cpython cimport *
cdef extern from "Python.h": cdef extern from "Python.h":
ctypedef char* const_char_ptr "const char*" ctypedef char* const_char_ptr "const char*"
ctypedef char* const_void_ptr "const void*"
ctypedef struct PyObject ctypedef struct PyObject
cdef int PyObject_AsReadBuffer(object o, const_void_ptr* buff, Py_ssize_t* buf_len) except -1
cdef object PyBytes_FromStringAndSize(const_char_ptr b, Py_ssize_t len) from libc.stdlib cimport *
cdef PyObject* Py_True from libc.string cimport *
cdef PyObject* Py_False
cdef object PyUnicode_AsUTF8String(object)
cdef long long PyLong_AsLongLong(object o)
cdef unsigned long long PyLong_AsUnsignedLongLong(object o)
cdef bint PyBool_Check(object o)
cdef bint PyDict_Check(object o)
cdef bint PySequence_Check(object o)
cdef bint PyLong_Check(object o)
cdef bint PyInt_Check(object o)
cdef bint PyFloat_Check(object o)
cdef bint PyBytes_Check(object o)
cdef bint PyUnicode_Check(object o)
cdef bint PyCallable_Check(object o)
cdef void Py_INCREF(object o)
cdef void Py_DECREF(object o)
cdef extern from "stdlib.h":
void* malloc(size_t)
void* realloc(void*, size_t)
void free(void*)
cdef extern from "string.h":
void* memcpy(char* dst, char* src, size_t size)
void* memmove(char* dst, char* src, size_t size)
cdef extern from "pack.h": cdef extern from "pack.h":
struct msgpack_packer: struct msgpack_packer:
@@ -104,10 +81,10 @@ cdef class Packer(object):
ret = msgpack_pack_false(&self.pk) ret = msgpack_pack_false(&self.pk)
elif PyLong_Check(o): elif PyLong_Check(o):
if o > 0: if o > 0:
ullval = PyLong_AsUnsignedLongLong(o) ullval = o
ret = msgpack_pack_unsigned_long_long(&self.pk, ullval) ret = msgpack_pack_unsigned_long_long(&self.pk, ullval)
else: else:
llval = PyLong_AsLongLong(o) llval = o
ret = msgpack_pack_long_long(&self.pk, llval) ret = msgpack_pack_long_long(&self.pk, llval)
elif PyInt_Check(o): elif PyInt_Check(o):
longval = o longval = o
@@ -160,7 +137,7 @@ cdef class Packer(object):
def pack(object o, object stream, default=None): def pack(object o, object stream, default=None):
"""pack an object `o` and write it to stream).""" """pack an object `o` and write it to stream)."""
packer = Packer(default) packer = Packer(default=default)
stream.write(packer.pack(o)) stream.write(packer.pack(o))
def packb(object o, default=None): def packb(object o, default=None):
@@ -189,12 +166,16 @@ cdef extern from "unpack.h":
object template_data(template_context* ctx) object template_data(template_context* ctx)
def unpackb(bytes packed_bytes, object object_hook=None, object list_hook=None): def unpackb(object packed, object object_hook=None, object list_hook=None):
"""Unpack packed_bytes to object. Returns an unpacked object.""" """Unpack packed_bytes to object. Returns an unpacked object."""
cdef const_char_ptr p = packed_bytes
cdef template_context ctx cdef template_context ctx
cdef size_t off = 0 cdef size_t off = 0
cdef int ret cdef int ret
cdef char* buf
cdef Py_ssize_t buf_len
PyObject_AsReadBuffer(packed, <const_void_ptr*>&buf, &buf_len)
template_init(&ctx) template_init(&ctx)
ctx.user.use_list = 0 ctx.user.use_list = 0
ctx.user.object_hook = ctx.user.list_hook = NULL ctx.user.object_hook = ctx.user.list_hook = NULL
@@ -206,7 +187,7 @@ def unpackb(bytes packed_bytes, object object_hook=None, object list_hook=None):
if not PyCallable_Check(list_hook): if not PyCallable_Check(list_hook):
raise TypeError("list_hook must be a callable.") raise TypeError("list_hook must be a callable.")
ctx.user.list_hook = <PyObject*>list_hook ctx.user.list_hook = <PyObject*>list_hook
ret = template_execute(&ctx, p, len(packed_bytes), &off) ret = template_execute(&ctx, buf, buf_len, &off)
if ret == 1: if ret == 1:
return template_data(&ctx) return template_data(&ctx)
else: else:
@@ -216,8 +197,8 @@ unpacks = unpackb
def unpack(object stream, object object_hook=None, object list_hook=None): def unpack(object stream, object object_hook=None, object list_hook=None):
"""unpack an object from stream.""" """unpack an object from stream."""
packed = stream.read() return unpackb(stream.read(),
return unpackb(packed, object_hook=object_hook, list_hook=list_hook) object_hook=object_hook, list_hook=list_hook)
cdef class UnpackIterator(object): cdef class UnpackIterator(object):
cdef object unpacker cdef object unpacker
@@ -232,21 +213,12 @@ cdef class UnpackIterator(object):
return self return self
cdef class Unpacker(object): cdef class Unpacker(object):
"""Unpacker(file_like=None, read_size=1024*1024) """Unpacker(read_size=1024*1024)
Streaming unpacker. Streaming unpacker.
file_like must have read(n) method.
read_size is used like file_like.read(read_size) read_size is used like file_like.read(read_size)
If file_like is None, you can ``feed()`` bytes. ``feed()`` is example:
useful for unpacking from non-blocking stream.
exsample 1:
unpacker = Unpacker(afile)
for o in unpacker:
do_something(o)
example 2:
unpacker = Unpacker() unpacker = Unpacker()
while 1: while 1:
buf = astream.read() buf = astream.read()
@@ -254,13 +226,11 @@ cdef class Unpacker(object):
for o in unpacker: for o in unpacker:
do_something(o) do_something(o)
""" """
cdef template_context ctx cdef template_context ctx
cdef char* buf cdef char* buf
cdef size_t buf_size, buf_head, buf_tail cdef size_t buf_size, buf_head, buf_tail
cdef object file_like cdef object file_like
cdef int read_size cdef int read_size
cdef object waiting_bytes
cdef bint use_list cdef bint use_list
cdef object object_hook cdef object object_hook
@@ -268,7 +238,6 @@ cdef class Unpacker(object):
self.buf = NULL self.buf = NULL
def __dealloc__(self): def __dealloc__(self):
if self.buf:
free(self.buf); free(self.buf);
def __init__(self, file_like=None, int read_size=0, bint use_list=0, def __init__(self, file_like=None, int read_size=0, bint use_list=0,
@@ -278,7 +247,6 @@ cdef class Unpacker(object):
self.use_list = use_list self.use_list = use_list
self.file_like = file_like self.file_like = file_like
self.read_size = read_size self.read_size = read_size
self.waiting_bytes = []
self.buf = <char*>malloc(read_size) self.buf = <char*>malloc(read_size)
self.buf_size = read_size self.buf_size = read_size
self.buf_head = 0 self.buf_head = 0
@@ -295,64 +263,48 @@ cdef class Unpacker(object):
raise TypeError("object_hook must be a callable.") raise TypeError("object_hook must be a callable.")
self.ctx.user.list_hook = <PyObject*>list_hook self.ctx.user.list_hook = <PyObject*>list_hook
def feed(self, bytes next_bytes): def feed(self, object next_bytes):
self.waiting_bytes.append(next_bytes) cdef char* buf
cdef Py_ssize_t buf_len
PyObject_AsReadBuffer(next_bytes, <const_void_ptr*>&buf, &buf_len)
self.append_buffer(buf, buf_len)
cdef append_buffer(self): cdef append_buffer(self, void* _buf, Py_ssize_t _buf_len):
cdef char* buf = self.buf cdef:
cdef Py_ssize_t tail = self.buf_tail char* buf = self.buf
cdef Py_ssize_t l size_t head = self.buf_head
cdef bytes b size_t tail = self.buf_tail
size_t buf_size = self.buf_size
size_t new_size
for b in self.waiting_bytes: if tail + _buf_len > buf_size:
l = len(b) if ((tail - head) + _buf_len)*2 < buf_size:
memcpy(buf + tail, <char*>(b), l)
tail += l
self.buf_tail = tail
del self.waiting_bytes[:]
# prepare self.buf
cdef fill_buffer(self):
cdef Py_ssize_t add_size
if self.file_like is not None:
next_bytes = self.file_like.read(self.read_size)
if next_bytes:
self.waiting_bytes.append(next_bytes)
else:
self.file_like = None
if not self.waiting_bytes:
return
add_size = 0
for b in self.waiting_bytes:
add_size += len(b)
cdef char* buf = self.buf
cdef size_t head = self.buf_head
cdef size_t tail = self.buf_tail
cdef size_t size = self.buf_size
if self.buf_tail + add_size <= self.buf_size:
# do nothing.
pass
if self.buf_tail - self.buf_head + add_size < self.buf_size:
# move to front. # move to front.
memmove(buf, buf + head, tail - head) memmove(buf, buf + head, tail - head)
tail -= head tail -= head
head = 0 head = 0
else: else:
# expand buffer # expand buffer.
size = tail + add_size new_size = tail + _buf_len
buf = <char*>realloc(<void*>buf, size) if new_size < buf_size*2:
new_size = buf_size*2
buf = <char*>realloc(buf, new_size)
buf_size = new_size
self.buf = buf memcpy(buf + tail, <char*>(_buf), _buf_len)
self.buf_head = head self.buf_head = head
self.buf_tail = tail self.buf_size = buf_size
self.buf_size = size self.buf_tail = tail + _buf_len
self.append_buffer() # prepare self.buf from file_like
cdef fill_buffer(self):
if self.file_like is not None:
next_bytes = self.file_like.read(self.read_size)
if next_bytes:
self.append_buffer(PyBytes_AsString(next_bytes),
PyBytes_Size(next_bytes))
else:
self.file_like = None
cpdef unpack(self): cpdef unpack(self):
"""unpack one object""" """unpack one object"""

View File

@@ -7,10 +7,10 @@ from msgpack import packb, unpackb
def test_unpack_buffer(): def test_unpack_buffer():
from array import array from array import array
buf = array('b') buf = array('c')
buf.fromstring(packb(['foo', 'bar'])) buf.fromstring(packb(('foo', 'bar')))
obj = unpackb(buf) obj = unpackb(buf)
assert_equal(['foo', 'bar'], obj) assert_equal(('foo', 'bar'), obj)
if __name__ == '__main__': if __name__ == '__main__':
main() main()

View File

@@ -1,22 +1,22 @@
#!/usr/bin/env python #!/usr/bin/env python
# coding: utf-8 # coding: utf-8
from __future__ import unicode_literals, print_function from __future__ import unicode_literals
from msgpack import Unpacker from msgpack import Unpacker
def test_foobar(): def test_foobar():
unpacker = Unpacker(read_size=3) unpacker = Unpacker(read_size=3)
unpacker.feed(b'foobar') unpacker.feed(b'foobar')
assert unpacker.unpack() == ord('f') assert unpacker.unpack() == ord(b'f')
assert unpacker.unpack() == ord('o') assert unpacker.unpack() == ord(b'o')
assert unpacker.unpack() == ord('o') assert unpacker.unpack() == ord(b'o')
assert unpacker.unpack() == ord('b') assert unpacker.unpack() == ord(b'b')
assert unpacker.unpack() == ord('a') assert unpacker.unpack() == ord(b'a')
assert unpacker.unpack() == ord('r') assert unpacker.unpack() == ord(b'r')
try: try:
o = unpacker.unpack() o = unpacker.unpack()
print("Oops!", o) print "Oops!", o
assert 0 assert 0
except StopIteration: except StopIteration:
assert 1 assert 1