Unpacker: add tell() (msgpack#227)

jfolz · methane · commit a8d9162ca6cf · 2017-04-30T02:33:20.000+09:00
diff --git 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
@@ -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
@@ -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()
