Unpacker: add tell() (#227)

3 files changed, 32 insertions, 0 deletions

diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
index a9801ee..dabc5f7 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 d2eb9f4..508fd06 100644
--- a/msgpack/fallback.py
+++ b/msgpack/fallback.py
@@ -244,6 +244,7 @@ class Unpacker(object):
         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 @@ class Unpacker(object):
 
     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 @@ class Unpacker(object):
         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 45f4cc7..59718f5 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()