fallback: add actual rollback and add a testcase for it

Signed-off-by: Bas Westerbaan <bas@westerbaan.name>

Author: bwesterb

msgpack/fallback.py

@@ -208,17 +208,13 @@ def _fb_got_extradata(self):
     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 @@ def _fb_read(self, n, write_bytes=None):
                 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 @@ def _fb_unpack(self, execute=EX_CONSTRUCT, write_bytes=None):
         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()

test/test_sequnpack.py

@@ -7,6 +7,21 @@
 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')