mirror of
https://github.com/msgpack/msgpack-python.git
synced 2025-10-19 12:03:15 +00:00
Fix Unpacker.tell() (#427)
Fixes #426. Co-authored-by: folz <joachim.folz@dfki.de>
This commit is contained in:
parent
b04690012d
commit
c1b1a23f62
3 changed files with 45 additions and 9 deletions
|
@ -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):
|
||||
|
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue