8000 Unpacker: add tell() (#227) · wumingcp-ps/msgpack-python@a8d9162 · GitHub
[go: up one dir, main page]

Skip to content

Commit a8d9162

Browse files
jfolzmethane
authored andcommitted
Unpacker: add tell() (msgpack#227)
1 parent 3388e4a commit a8d9162

File tree

3 files changed

+32
-0
lines changed

3 files changed

+32
-0
lines changed

msgpack/_unpacker.pyx

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change< 8000 /th>
@@ -29,6 +29,7 @@ cdef extern from "Python.h":
2929
from libc.stdlib cimport *
3030
from libc.string cimport *
3131
from libc.limits cimport *
32+
ctypedef unsigned long long uint64_t
3233

3334
from msgpack.exceptions import (
3435
BufferFull,
@@ -314,6 +315,7 @@ cdef class Unpacker(object):
314315
cdef object object_hook, object_pairs_hook, list_hook, ext_hook
315316
cdef object encoding, unicode_errors
316317
cdef Py_ssize_t max_buffer_size
318+
cdef uint64_t stream_offset
317319

318320
def __cinit__(self):
319321
self.buf = NULL
@@ -358,6 +360,7 @@ cdef class Unpacker(object):
358360
self.buf_size = read_size
359361
self.buf_head = 0
360362
self.buf_tail = 0
363+
self.stream_offset = 0
361364

362365
if encoding is not None:
363366
if isinstance(encoding, unicode):
@@ -468,6 +471,7 @@ cdef class Unpacker(object):
468471

469472
try:
470473
ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
474+
self.stream_offset += self.buf_head - prev_head
471475
if write_bytes is not None:
472476
write_bytes(PyBytes_FromStringAndSize(self.buf + prev_head, self.buf_head - prev_head))
473477

@@ -534,6 +538,9 @@ cdef class Unpacker(object):
534538
"""
535539
return self._unpack(read_map_header, write_bytes)
536540

541+
def tell(self):
542+
return self.stream_offset
543+
537544
def __iter__(self):
538545
return self
539546

msgpack/fallback.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,7 @@ def __init__(self, file_like=None, read_size=0, use_list=True,
244244
self._max_array_len = max_array_len
245245
self._max_map_len = max_map_len
246246
self._max_ext_len = max_ext_len
247+
self._stream_offset = 0
247248

248249
if list_hook is not None and not callable(list_hook):
249250
raise TypeError('`list_hook` is not callable')
@@ -266,6 +267,7 @@ def feed(self, next_bytes):
266267

267268
def _consume(self):
268269
""" Gets rid of the used parts of the buffer. """
270+
self._stream_offset += self._buff_i - self._buf_checkpoint
269271
self._buf_checkpoint = self._buff_i
270272

271273
def _got_extradata(self):
@@ -629,6 +631,9 @@ def read_map_header(self, write_bytes=None):
629631
self._consume()
630632
return ret
631633

634+
def tell(self):
635+
return self._stream_offset
636+
632637

633638
class Packer(object):
634639
"""

test/test_sequnpack.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import io
55
from msgpack import Unpacker, BufferFull
6+
from msgpack import pack
67
from msgpack.exceptions import OutOfData
78
from pytest import raises
89

@@ -96,3 +97,22 @@ def test_issue124():
9697
unpacker.feed(b"!")
9798
assert tuple(unpacker) == (b'!',)
9899
assert tuple(unpacker) == ()
100+
101+
102+
def test_unpack_tell():
103+
stream = io.BytesIO()
104+
messages = [2**i-1 for i in range(65)]
105+
messages += [-(2**i) for i in range(1, 64)]
106+
messages += [b'hello', b'hello'*1000, list(range(20)),
107+
{i: bytes(i)*i for i in range(10)},
108+
{i: bytes(i)*i for i in range(32)}]
109+
offsets = []
110+
for m in messages:
111+
pack(m, stream)
112+
offsets.append(stream.tell())
113+
stream.seek(0)
114+
unpacker = Unpacker(stream)
115+
for m, o in zip(messages, offsets):
116+
m2 = next(unpacker)
117+
assert m == m2
118+
assert o == unpacker.tell()

0 commit comments

Comments
 (0)
0