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