Merge pull request msgpack#87 from msgpack/fix-83

methane · methane · commit 7b24d0fe5a20 · 2014-02-13T03:24:09.000+09:00
Feed data from file before _unpack()
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
@@ -327,8 +327,18 @@ cdef class Unpacker(object):
         cdef int ret
         cdef object obj
         cdef size_t prev_head
+
+        if self.buf_head >= self.buf_tail and self.file_like is not None:
+            self.read_from_file()
+
         while 1:
             prev_head = self.buf_head
+            if prev_head >= self.buf_tail:
+                if iter:
+                    raise StopIteration("No more data to unpack.")
+                else:
+                    raise OutOfData("No more data to unpack.")
+
             ret = execute(&self.ctx, self.buf, self.buf_tail, &self.buf_head)
             if write_bytes is not None:
                 write_bytes(PyBytes_FromStringAndSize(self.buf + prev_head, self.buf_head - prev_head))
diff --git a/test/test_unpack_file.py b/test/test_unpack_file.py
@@ -0,0 +1,19 @@
+from io import BytesIO
+from msgpack import Unpacker, packb, OutOfData
+from pytest import raises
+
+
+def test_unpack_array_header_from_file():
+    f = BytesIO(packb([1,2,3,4]))
+    unpacker = Unpacker(f)
+    assert unpacker.read_array_header() == 4
+    assert unpacker.unpack() == 1
+    assert unpacker.unpack() == 2
+    assert unpacker.unpack() == 3
+    assert unpacker.unpack() == 4
+    with raises(OutOfData):
+        unpacker.unpack()
+
+
+if __name__ == '__main__':
+    test_unpack_array_header_from_file()
