_unpacker.pyx 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426
  1. # coding: utf-8
  2. #cython: embedsignature=True
  3. from cpython cimport *
  4. cdef extern from "Python.h":
  5. ctypedef struct PyObject
  6. cdef int PyObject_AsReadBuffer(object o, const void** buff, Py_ssize_t* buf_len) except -1
  7. from libc.stdlib cimport *
  8. from libc.string cimport *
  9. from libc.limits cimport *
  10. from msgpack.exceptions import (
  11. BufferFull,
  12. OutOfData,
  13. UnpackValueError,
  14. ExtraData,
  15. )
  16. from msgpack import ExtType
  17. cdef extern from "unpack.h":
  18. ctypedef struct msgpack_user:
  19. bint use_list
  20. PyObject* object_hook
  21. bint has_pairs_hook # call object_hook with k-v pairs
  22. PyObject* list_hook
  23. PyObject* ext_hook
  24. char *encoding
  25. char *unicode_errors
  26. ctypedef struct unpack_context:
  27. msgpack_user user
  28. PyObject* obj
  29. size_t count
  30. ctypedef int (*execute_fn)(unpack_context* ctx, const char* data,
  31. size_t len, size_t* off) except? -1
  32. execute_fn unpack_construct
  33. execute_fn unpack_skip
  34. execute_fn read_array_header
  35. execute_fn read_map_header
  36. void unpack_init(unpack_context* ctx)
  37. object unpack_data(unpack_context* ctx)
  38. cdef inline init_ctx(unpack_context *ctx,
  39. object object_hook, object object_pairs_hook,
  40. object list_hook, object ext_hook,
  41. bint use_list, char* encoding, char* unicode_errors):
  42. unpack_init(ctx)
  43. ctx.user.use_list = use_list
  44. ctx.user.object_hook = ctx.user.list_hook = <PyObject*>NULL
  45. if object_hook is not None and object_pairs_hook is not None:
  46. raise TypeError("object_pairs_hook and object_hook are mutually exclusive.")
  47. if object_hook is not None:
  48. if not PyCallable_Check(object_hook):
  49. raise TypeError("object_hook must be a callable.")
  50. ctx.user.object_hook = <PyObject*>object_hook
  51. if object_pairs_hook is None:
  52. ctx.user.has_pairs_hook = False
  53. else:
  54. if not PyCallable_Check(object_pairs_hook):
  55. raise TypeError("object_pairs_hook must be a callable.")
  56. ctx.user.object_hook = <PyObject*>object_pairs_hook
  57. ctx.user.has_pairs_hook = True
  58. if list_hook is not None:
  59. if not PyCallable_Check(list_hook):
  60. raise TypeError("list_hook must be a callable.")
  61. ctx.user.list_hook = <PyObject*>list_hook
  62. if ext_hook is not None:
  63. if not PyCallable_Check(ext_hook):
  64. raise TypeError("ext_hook must be a callable.")
  65. ctx.user.ext_hook = <PyObject*>ext_hook
  66. ctx.user.encoding = encoding
  67. ctx.user.unicode_errors = unicode_errors
  68. def default_read_extended_type(typecode, data):
  69. raise NotImplementedError("Cannot decode extended type with typecode=%d" % typecode)
  70. def unpackb(object packed, object object_hook=None, object list_hook=None,
  71. bint use_list=1, encoding=None, unicode_errors="strict",
  72. object_pairs_hook=None, ext_hook=ExtType):
  73. """
  74. Unpack packed_bytes to object. Returns an unpacked object.
  75. Raises `ValueError` when `packed` contains extra bytes.
  76. See :class:`Unpacker` for options.
  77. """
  78. cdef unpack_context ctx
  79. cdef size_t off = 0
  80. cdef int ret
  81. cdef char* buf
  82. cdef Py_ssize_t buf_len
  83. cdef char* cenc = NULL
  84. cdef char* cerr = NULL
  85. PyObject_AsReadBuffer(packed, <const void**>&buf, &buf_len)
  86. if encoding is not None:
  87. if isinstance(encoding, unicode):
  88. encoding = encoding.encode('ascii')
  89. cenc = PyBytes_AsString(encoding)
  90. if unicode_errors is not None:
  91. if isinstance(unicode_errors, unicode):
  92. unicode_errors = unicode_errors.encode('ascii')
  93. cerr = PyBytes_AsString(unicode_errors)
  94. init_ctx(&ctx, object_hook, object_pairs_hook, list_hook, ext_hook,
  95. use_list, cenc, cerr)
  96. ret = unpack_construct(&ctx, buf, buf_len, &off)
  97. if ret == 1:
  98. obj = unpack_data(&ctx)
  99. if off < buf_len:
  100. raise ExtraData(obj, PyBytes_FromStringAndSize(buf+off, buf_len-off))
  101. return obj
  102. else:
  103. raise UnpackValueError("Unpack failed: error = %d" % (ret,))
  104. def unpack(object stream, object object_hook=None, object list_hook=None,
  105. bint use_list=1, encoding=None, unicode_errors="strict",
  106. object_pairs_hook=None,
  107. ):
  108. """
  109. Unpack an object from `stream`.
  110. Raises `ValueError` when `stream` has extra bytes.
  111. See :class:`Unpacker` for options.
  112. """
  113. return unpackb(stream.read(), use_list=use_list,
  114. object_hook=object_hook, object_pairs_hook=object_pairs_hook, list_hook=list_hook,
  115. encoding=encoding, unicode_errors=unicode_errors,
  116. )
  117. cdef class Unpacker(object):
  118. """
  119. Streaming unpacker.
  120. arguments:
  121. :param file_like:
  122. File-like object having `.read(n)` method.
  123. If specified, unpacker reads serialized data from it and :meth:`feed()` is not usable.
  124. :param int read_size:
  125. Used as `file_like.read(read_size)`. (default: `min(1024**2, max_buffer_size)`)
  126. :param bool use_list:
  127. If true, unpack msgpack array to Python list.
  128. Otherwise, unpack to Python tuple. (default: True)
  129. :param callable object_hook:
  130. When specified, it should be callable.
  131. Unpacker calls it with a dict argument after unpacking msgpack map.
  132. (See also simplejson)
  133. :param callable object_pairs_hook:
  134. When specified, it should be callable.
  135. Unpacker calls it with a list of key-value pairs after unpacking msgpack map.
  136. (See also simplejson)
  137. :param str encoding:
  138. Encoding used for decoding msgpack raw.
  139. If it is None (default), msgpack raw is deserialized to Python bytes.
  140. :param str unicode_errors:
  141. Used for decoding msgpack raw with *encoding*.
  142. (default: `'strict'`)
  143. :param int max_buffer_size:
  144. Limits size of data waiting unpacked. 0 means system's INT_MAX (default).
  145. Raises `BufferFull` exception when it is insufficient.
  146. You shoud set this parameter when unpacking data from untrasted source.
  147. example of streaming deserialize from file-like object::
  148. unpacker = Unpacker(file_like)
  149. for o in unpacker:
  150. process(o)
  151. example of streaming deserialize from socket::
  152. unpacker = Unpacker()
  153. while True:
  154. buf = sock.recv(1024**2)
  155. if not buf:
  156. break
  157. unpacker.feed(buf)
  158. for o in unpacker:
  159. process(o)
  160. """
  161. cdef unpack_context ctx
  162. cdef char* buf
  163. cdef size_t buf_size, buf_head, buf_tail
  164. cdef object file_like
  165. cdef object file_like_read
  166. cdef Py_ssize_t read_size
  167. # To maintain refcnt.
  168. cdef object object_hook, object_pairs_hook, list_hook, ext_hook
  169. cdef object encoding, unicode_errors
  170. cdef size_t max_buffer_size
  171. def __cinit__(self):
  172. self.buf = NULL
  173. def __dealloc__(self):
  174. free(self.buf)
  175. self.buf = NULL
  176. def __init__(self, file_like=None, Py_ssize_t read_size=0, bint use_list=1,
  177. object object_hook=None, object object_pairs_hook=None, object list_hook=None,
  178. str encoding=None, str unicode_errors='strict', int max_buffer_size=0,
  179. object ext_hook=ExtType):
  180. cdef char *cenc=NULL,
  181. cdef char *cerr=NULL
  182. self.object_hook = object_hook
  183. self.object_pairs_hook = object_pairs_hook
  184. self.list_hook = list_hook
  185. self.ext_hook = ext_hook
  186. self.file_like = file_like
  187. if file_like:
  188. self.file_like_read = file_like.read
  189. if not PyCallable_Check(self.file_like_read):
  190. raise TypeError("`file_like.read` must be a callable.")
  191. if not max_buffer_size:
  192. max_buffer_size = INT_MAX
  193. if read_size > max_buffer_size:
  194. raise ValueError("read_size should be less or equal to max_buffer_size")
  195. if not read_size:
  196. read_size = min(max_buffer_size, 1024**2)
  197. self.max_buffer_size = max_buffer_size
  198. self.read_size = read_size
  199. self.buf = <char*>malloc(read_size)
  200. if self.buf == NULL:
  201. raise MemoryError("Unable to allocate internal buffer.")
  202. self.buf_size = read_size
  203. self.buf_head = 0
  204. self.buf_tail = 0
  205. if encoding is not None:
  206. if isinstance(encoding, unicode):
  207. self.encoding = encoding.encode('ascii')
  208. else:
  209. self.encoding = encoding
  210. cenc = PyBytes_AsString(self.encoding)
  211. if unicode_errors is not None:
  212. if isinstance(unicode_errors, unicode):
  213. self.unicode_errors = unicode_errors.encode('ascii')
  214. else:
  215. self.unicode_errors = unicode_errors
  216. cerr = PyBytes_AsString(self.unicode_errors)
  217. init_ctx(&self.ctx, object_hook, object_pairs_hook, list_hook,
  218. ext_hook, use_list, cenc, cerr)
  219. def feed(self, object next_bytes):
  220. """Append `next_bytes` to internal buffer."""
  221. cdef Py_buffer pybuff
  222. if self.file_like is not None:
  223. raise AssertionError(
  224. "unpacker.feed() is not be able to use with `file_like`.")
  225. PyObject_GetBuffer(next_bytes, &pybuff, PyBUF_SIMPLE)
  226. try:
  227. self.append_buffer(<char*>pybuff.buf, pybuff.len)
  228. finally:
  229. PyBuffer_Release(&pybuff)
  230. cdef append_buffer(self, void* _buf, Py_ssize_t _buf_len):
  231. cdef:
  232. char* buf = self.buf
  233. char* new_buf
  234. size_t head = self.buf_head
  235. size_t tail = self.buf_tail
  236. size_t buf_size = self.buf_size
  237. size_t new_size
  238. if tail + _buf_len > buf_size:
  239. if ((tail - head) + _buf_len) <= buf_size:
  240. # move to front.
  241. memmove(buf, buf + head, tail - head)
  242. tail -= head
  243. head = 0
  244. else:
  245. # expand buffer.
  246. new_size = (tail-head) + _buf_len
  247. if new_size > self.max_buffer_size:
  248. raise BufferFull
  249. new_size = min(new_size*2, self.max_buffer_size)
  250. new_buf = <char*>malloc(new_size)
  251. if new_buf == NULL:
  252. # self.buf still holds old buffer and will be freed during
  253. # obj destruction
  254. raise MemoryError("Unable to enlarge internal buffer.")
  255. memcpy(new_buf, buf + head, tail - head)
  256. free(buf)
  257. buf = new_buf
  258. buf_size = new_size
  259. tail -= head
  260. head = 0
  261. memcpy(buf + tail, <char*>(_buf), _buf_len)
  262. self.buf = buf
  263. self.buf_head = head
  264. self.buf_size = buf_size
  265. self.buf_tail = tail + _buf_len
  266. cdef read_from_file(self):
  267. next_bytes = self.file_like_read(
  268. min(self.read_size,
  269. self.max_buffer_size - (self.buf_tail - self.buf_head)
  270. ))
  271. if next_bytes:
  272. self.append_buffer(PyBytes_AsString(next_bytes), PyBytes_Size(next_bytes))
  273. else:
  274. self.file_like = None
  275. cdef object _unpack(self, execute_fn execute, object write_bytes, bint iter=0):
  276. cdef int ret
  277. cdef object obj
  278. cdef size_t prev_head
  279. if self.buf_head >= self.buf_tail and self.file_like is not None:
  280. self.read_from_file()
  281. while 1:
  282. prev_head = self.buf_head
  283. if prev_head >= self.buf_tail:
  284. if iter:
  285. raise StopIteration("No more data to unpack.")
  286. else:
  287. raise OutOfData("No more data to unpack.")
  288. ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
  289. if write_bytes is not None:
  290. write_bytes(PyBytes_FromStringAndSize(self.buf + prev_head, self.buf_head - prev_head))
  291. if ret == 1:
  292. obj = unpack_data(&self.ctx)
  293. unpack_init(&self.ctx)
  294. return obj
  295. elif ret == 0:
  296. if self.file_like is not None:
  297. self.read_from_file()
  298. continue
  299. if iter:
  300. raise StopIteration("No more data to unpack.")
  301. else:
  302. raise OutOfData("No more data to unpack.")
  303. else:
  304. raise ValueError("Unpack failed: error = %d" % (ret,))
  305. def read_bytes(self, Py_ssize_t nbytes):
  306. """read a specified number of raw bytes from the stream"""
  307. cdef size_t nread
  308. nread = min(self.buf_tail - self.buf_head, nbytes)
  309. ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
  310. self.buf_head += nread
  311. if len(ret) < nbytes and self.file_like is not None:
  312. ret += self.file_like.read(nbytes - len(ret))
  313. return ret
  314. def unpack(self, object write_bytes=None):
  315. """
  316. unpack one object
  317. If write_bytes is not None, it will be called with parts of the raw
  318. message as it is unpacked.
  319. Raises `OutOfData` when there are no more bytes to unpack.
  320. """
  321. return self._unpack(unpack_construct, write_bytes)
  322. def skip(self, object write_bytes=None):
  323. """
  324. read and ignore one object, returning None
  325. If write_bytes is not None, it will be called with parts of the raw
  326. message as it is unpacked.
  327. Raises `OutOfData` when there are no more bytes to unpack.
  328. """
  329. return self._unpack(unpack_skip, write_bytes)
  330. def read_array_header(self, object write_bytes=None):
  331. """assuming the next object is an array, return its size n, such that
  332. the next n unpack() calls will iterate over its contents.
  333. Raises `OutOfData` when there are no more bytes to unpack.
  334. """
  335. return self._unpack(read_array_header, write_bytes)
  336. def read_map_header(self, object write_bytes=None):
  337. """assuming the next object is a map, return its size n, such that the
  338. next n * 2 unpack() calls will iterate over its key-value pairs.
  339. Raises `OutOfData` when there are no more bytes to unpack.
  340. """
  341. return self._unpack(read_map_header, write_bytes)
  342. def __iter__(self):
  343. return self
  344. def __next__(self):
  345. return self._unpack(unpack_construct, None, 1)
  346. # for debug.
  347. #def _buf(self):
  348. # return PyString_FromStringAndSize(self.buf, self.buf_tail)
  349. #def _off(self):
  350. # return self.buf_head