diff options
author | jfolz <theriddling@gmail.com> | 2020-06-08 05:14:50 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-06-08 12:14:50 +0900 |
commit | c1b1a23f62d5e0ec39a1910d2e9580ce1c13a1cf (patch) | |
tree | 9726bba0edf80ba9515495c315d7966f7b7aa37e | |
parent | b04690012d5d77cfe5074893686c4d55ec780300 (diff) | |
download | msgpack-python-c1b1a23f62d5e0ec39a1910d2e9580ce1c13a1cf.tar.gz |
Fix Unpacker.tell() (#427)
Fixes #426.
Co-authored-by: folz <joachim.folz@dfki.de>
-rw-r--r-- | msgpack/_unpacker.pyx | 10 | ||||
-rw-r--r-- | msgpack/fallback.py | 15 | ||||
-rw-r--r-- | test/test_unpack.py | 29 |
3 files changed, 45 insertions, 9 deletions
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx index 43c93a2..4340e04 100644 --- a/msgpack/_unpacker.pyx +++ b/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): diff --git a/msgpack/fallback.py b/msgpack/fallback.py index 9f6665b..1e0bbe9 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -365,18 +365,19 @@ class Unpacker(object): 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 @@ class Unpacker(object): 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 diff --git a/test/test_unpack.py b/test/test_unpack.py index bc74c4d..057b7bf 100644 --- a/test/test_unpack.py +++ b/test/test_unpack.py @@ -3,6 +3,11 @@ import sys 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 test_unpacker_ext_hook(): 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() |