Fix Unpacker.tell() (#427)

Fixes #426.

Co-authored-by: folz <joachim.folz@dfki.de>
This commit is contained in:
jfolz 2020-06-08 05:14:50 +02:00 committed by GitHub
parent b04690012d
commit c1b1a23f62
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 45 additions and 9 deletions

View file

@ -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):

View file

@ -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

View file

@ -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()