From d2f549a47094b2a29cc94bc50029ebdc85508861 Mon Sep 17 00:00:00 2001 From: Bas Westerbaan Date: Tue, 29 Jan 2013 02:58:26 +0100 Subject: fallback: add actual rollback and add a testcase for it Signed-off-by: Bas Westerbaan --- msgpack/fallback.py | 21 +++++++++++++-------- test/test_sequnpack.py | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/msgpack/fallback.py b/msgpack/fallback.py index a866ff1..2f83a20 100644 --- a/msgpack/fallback.py +++ b/msgpack/fallback.py @@ -208,17 +208,13 @@ class Unpacker(object): def __iter__(self): return self - def next(self): - try: - ret = self._fb_unpack(EX_CONSTRUCT, None) - self._fb_consume() - return ret - except OutOfData: - raise StopIteration - def read_bytes(self, n): return self._fb_read(n) + def _fb_rollback(self): + self._fb_buf_i = 0 + self._fb_buf_o = 0 + def _fb_get_extradata(self): bufs = self._fb_buffers[self._fb_buf_i:] if bufs: @@ -244,6 +240,7 @@ class Unpacker(object): self._fb_buf_o = 0 self._fb_buf_i += 1 if len(ret) != n: + self._fb_rollback() raise OutOfData if write_bytes is not None: write_bytes(ret) @@ -363,6 +360,14 @@ class Unpacker(object): assert typ == TYPE_IMMEDIATE return obj + def next(self): + try: + ret = self._fb_unpack(EX_CONSTRUCT, None) + self._fb_consume() + return ret + except OutOfData: + raise StopIteration + def skip(self, write_bytes=None): self._fb_unpack(EX_SKIP, write_bytes) self._fb_consume() diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py index fc1f712..1da383c 100644 --- a/test/test_sequnpack.py +++ b/test/test_sequnpack.py @@ -7,6 +7,21 @@ from msgpack.exceptions import OutOfData from pytest import raises +def test_partialdata(): + unpacker = Unpacker() + unpacker.feed(b'\xa5') + with raises(StopIteration): next(iter(unpacker)) + unpacker.feed(b'h') + with raises(StopIteration): next(iter(unpacker)) + unpacker.feed(b'a') + with raises(StopIteration): next(iter(unpacker)) + unpacker.feed(b'l') + with raises(StopIteration): next(iter(unpacker)) + unpacker.feed(b'l') + with raises(StopIteration): next(iter(unpacker)) + unpacker.feed(b'o') + assert next(iter(unpacker)) == 'hallo' + def test_foobar(): unpacker = Unpacker(read_size=3, use_list=1) unpacker.feed(b'foobar') -- cgit v1.2.1