diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index a9801eef..dabc5f70 100644 --- a/msgpack/_unpacker.pyx +++ b/msgpack/_unpacker.pyx @@ -29,6 +29,7 @@ cdef extern from "Python.h": from libc.stdlib cimport * from libc.string cimport * from libc.limits cimport * +ctypedef unsigned long long uint64_t from msgpack.exceptions import ( BufferFull, @@ -314,6 +315,7 @@ cdef class Unpacker(object): cdef object object_hook, object_pairs_hook, list_hook, ext_hook cdef object encoding, unicode_errors cdef Py_ssize_t max_buffer_size + cdef uint64_t stream_offset def __cinit__(self): self.buf = NULL @@ -358,6 +360,7 @@ cdef class Unpacker(object): self.buf_size = read_size self.buf_head = 0 self.buf_tail = 0 + self.stream_offset = 0 if encoding is not None: if isinstance(encoding, unicode): @@ -468,6 +471,7 @@ cdef class Unpacker(object): try: ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head) + self.stream_offset += self.buf_head - prev_head if write_bytes is not None: write_bytes(PyBytes_FromStringAndSize(self.buf + prev_head, self.buf_head - prev_head)) @@ -534,6 +538,9 @@ cdef class Unpacker(object): """ return self._unpack(read_map_header, write_bytes) + def tell(self): + return self.stream_offset + def __iter__(self): return self diff --git a/msgpack/fallback.py b/msgpack/fallback.py index d2eb9f44..508fd06f 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -244,6 +244,7 @@ def __init__(self, file_like=None, read_size=0, use_list=True, self._max_array_len = max_array_len self._max_map_len = max_map_len self._max_ext_len = max_ext_len + self._stream_offset = 0 if list_hook is not None and not callable(list_hook): raise TypeError('`list_hook` is not callable') @@ -266,6 +267,7 @@ def feed(self, next_bytes): def _consume(self): """ Gets rid of the used parts of the buffer. """ + self._stream_offset += self._buff_i - self._buf_checkpoint self._buf_checkpoint = self._buff_i def _got_extradata(self): @@ -629,6 +631,9 @@ def read_map_header(self, write_bytes=None): self._consume() return ret + def tell(self): + return self._stream_offset + class Packer(object): """ diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py index 45f4cc78..59718f56 100644 --- a/test/test_sequnpack.py +++ b/test/test_sequnpack.py @@ -3,6 +3,7 @@ import io from msgpack import Unpacker, BufferFull +from msgpack import pack from msgpack.exceptions import OutOfData from pytest import raises @@ -96,3 +97,22 @@ def test_issue124(): unpacker.feed(b"!") assert tuple(unpacker) == (b'!',) assert tuple(unpacker) == () + + +def test_unpack_tell(): + stream = io.BytesIO() + messages = [2**i-1 for i in range(65)] + messages += [-(2**i) for i in range(1, 64)] + messages += [b'hello', b'hello'*1000, list(range(20)), + {i: bytes(i)*i for i in range(10)}, + {i: bytes(i)*i for i in range(32)}] + offsets = [] + for m in messages: + pack(m, stream) + offsets.append(stream.tell()) + stream.seek(0) + unpacker = Unpacker(stream) + for m, o in zip(messages, offsets): + m2 = next(unpacker) + assert m == m2 + assert o == unpacker.tell()