Fix Unpacker.tell() (#427)

Fixes #426.

Co-authored-by: folz <[email protected]>

Author: jfolz

msgpack/_unpacker.pyx

@@ -484,8 +484,10 @@ cdef class Unpacker(object):
         nread = min(self.buf_tail - self.buf_head, nbytes)
         ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
         self.buf_head += nread
-        if len(ret) < nbytes and self.file_like is not None:
-            ret += self.file_like.read(nbytes - len(ret))
+        if nread < nbytes and self.file_like is not None:
+            ret += self.file_like.read(nbytes - nread)
+            nread = len(ret)
+        self.stream_offset += nread
         return ret
 
     def unpack(self):
@@ -519,6 +521,10 @@ cdef class Unpacker(object):
         return self._unpack(read_map_header)
 
     def tell(self):
+        """Returns the current position of the Unpacker in bytes, i.e., the
+        number of bytes that were read from the input, also the starting
+        position of the next object.
+        """
         return self.stream_offset
 
     def __iter__(self):

msgpack/fallback.py

@@ -365,18 +365,19 @@ def _get_extradata(self):
         return self._buffer[self._buff_i :]
 
     def read_bytes(self, n):
-        ret = self._read(n)
+        ret = self._read(n, raise_outofdata=False)
         self._consume()
         return ret
 
-    def _read(self, n):
+    def _read(self, n, raise_outofdata=True):
         # (int) -> bytearray
-        self._reserve(n)
+        self._reserve(n, raise_outofdata=raise_outofdata)
         i = self._buff_i
-        self._buff_i = i + n
-        return self._buffer[i : i + n]
+        ret = self._buffer[i : i + n]
+        self._buff_i = i + len(ret)
+        return ret
 
-    def _reserve(self, n):
+    def _reserve(self, n, raise_outofdata=True):
         remain_bytes = len(self._buffer) - self._buff_i - n
 
         # Fast path: buffer has n bytes already
@@ -404,7 +405,7 @@ def _reserve(self, n):
             self._buffer += read_data
             remain_bytes -= len(read_data)
 
-        if len(self._buffer) < n + self._buff_i:
+        if len(self._buffer) < n + self._buff_i and raise_outofdata:
             self._buff_i = 0  # rollback
             raise OutOfData
 

test/test_unpack.py

@@ -3,6 +3,11 @@
 from msgpack import Unpacker, packb, OutOfData, ExtType
 from pytest import raises, mark
 
+try:
+    from itertools import izip as zip
+except ImportError:
+    pass
+
 
 def test_unpack_array_header_from_file():
     f = BytesIO(packb([1, 2, 3, 4]))
@@ -64,7 +69,31 @@ def _hook(self, code, data):
     assert unpacker.unpack() == {"a": ExtType(2, b"321")}
 
 
+def test_unpacker_tell():
+    objects = 1, 2, u"abc", u"def", u"ghi"
+    packed = b"\x01\x02\xa3abc\xa3def\xa3ghi"
+    positions = 1, 2, 6, 10, 14
+    unpacker = Unpacker(BytesIO(packed))
+    for obj, unp, pos in zip(objects, unpacker, positions):
+        assert obj == unp
+        assert pos == unpacker.tell()
+
+
+def test_unpacker_tell_read_bytes():
+    objects = 1, u"abc", u"ghi"
+    packed = b"\x01\x02\xa3abc\xa3def\xa3ghi"
+    raw_data = b"\x02", b"\xa3def", b""
+    lenghts = 1, 4, 999
+    positions = 1, 6, 14
+    unpacker = Unpacker(BytesIO(packed))
+    for obj, unp, pos, n, raw in zip(objects, unpacker, positions, lenghts, raw_data):
+        assert obj == unp
+        assert pos == unpacker.tell()
+        assert unpacker.read_bytes(n) == raw
+
+
 if __name__ == "__main__":
     test_unpack_array_header_from_file()
     test_unpacker_hook_refcnt()
     test_unpacker_ext_hook()
+    test_unpacker_tell()