mirror of
https://github.com/msgpack/msgpack-python.git
synced 2025-10-19 20:03:16 +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)
|
nread = min(self.buf_tail - self.buf_head, nbytes)
|
||||||
ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
|
ret = PyBytes_FromStringAndSize(self.buf + self.buf_head, nread)
|
||||||
self.buf_head += nread
|
self.buf_head += nread
|
||||||
if len(ret) < nbytes and self.file_like is not None:
|
if nread < nbytes and self.file_like is not None:
|
||||||
ret += self.file_like.read(nbytes - len(ret))
|
ret += self.file_like.read(nbytes - nread)
|
||||||
|
nread = len(ret)
|
||||||
|
self.stream_offset += nread
|
||||||
return ret
|
return ret
|
||||||
|
|
||||||
def unpack(self):
|
def unpack(self):
|
||||||
|
@ -519,6 +521,10 @@ cdef class Unpacker(object):
|
||||||
return self._unpack(read_map_header)
|
return self._unpack(read_map_header)
|
||||||
|
|
||||||
def tell(self):
|
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
|
return self.stream_offset
|
||||||
|
|
||||||
def __iter__(self):
|
def __iter__(self):
|
||||||
|
|
|
@ -365,18 +365,19 @@ class Unpacker(object):
|
||||||
return self._buffer[self._buff_i :]
|
return self._buffer[self._buff_i :]
|
||||||
|
|
||||||
def read_bytes(self, n):
|
def read_bytes(self, n):
|
||||||
ret = self._read(n)
|
ret = self._read(n, raise_outofdata=False)
|
||||||
self._consume()
|
self._consume()
|
||||||
return ret
|
return ret
|
||||||
|
|
||||||
def _read(self, n):
|
def _read(self, n, raise_outofdata=True):
|
||||||
# (int) -> bytearray
|
# (int) -> bytearray
|
||||||
self._reserve(n)
|
self._reserve(n, raise_outofdata=raise_outofdata)
|
||||||
i = self._buff_i
|
i = self._buff_i
|
||||||
self._buff_i = i + n
|
ret = self._buffer[i : i + n]
|
||||||
return 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
|
remain_bytes = len(self._buffer) - self._buff_i - n
|
||||||
|
|
||||||
# Fast path: buffer has n bytes already
|
# Fast path: buffer has n bytes already
|
||||||
|
@ -404,7 +405,7 @@ class Unpacker(object):
|
||||||
self._buffer += read_data
|
self._buffer += read_data
|
||||||
remain_bytes -= len(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
|
self._buff_i = 0 # rollback
|
||||||
raise OutOfData
|
raise OutOfData
|
||||||
|
|
||||||
|
|
|
@ -3,6 +3,11 @@ import sys
|
||||||
from msgpack import Unpacker, packb, OutOfData, ExtType
|
from msgpack import Unpacker, packb, OutOfData, ExtType
|
||||||
from pytest import raises, mark
|
from pytest import raises, mark
|
||||||
|
|
||||||
|
try:
|
||||||
|
from itertools import izip as zip
|
||||||
|
except ImportError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
def test_unpack_array_header_from_file():
|
def test_unpack_array_header_from_file():
|
||||||
f = BytesIO(packb([1, 2, 3, 4]))
|
f = BytesIO(packb([1, 2, 3, 4]))
|
||||||
|
@ -64,7 +69,31 @@ def test_unpacker_ext_hook():
|
||||||
assert unpacker.unpack() == {"a": ExtType(2, b"321")}
|
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__":
|
if __name__ == "__main__":
|
||||||
test_unpack_array_header_from_file()
|
test_unpack_array_header_from_file()
|
||||||
test_unpacker_hook_refcnt()
|
test_unpacker_hook_refcnt()
|
||||||
test_unpacker_ext_hook()
|
test_unpacker_ext_hook()
|
||||||
|
test_unpacker_tell()
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue