implement unpack_one also for the cython version, and add a test for it

antocuni · antocuni · commit ff858387d37d · 2013-10-19T01:49:03.000+02:00
diff --git a/msgpack/_unpacker.pyx b/msgpack/_unpacker.pyx
@@ -359,6 +359,24 @@ cdef class Unpacker(object):
         """
         return self._unpack(unpack_construct, write_bytes)
 
+    def unpack_one(self, object write_bytes=None):
+        """
+        unpack one object
+
+        If write_bytes is not None, it will be called with parts of the raw
+        message as it is unpacked.
+
+        Raises `UnpackValueError` if there are no more bytes to unpack.
+        Raises ``ExtraData`` if there are still bytes left after the unpacking.
+        """
+        try:
+            result = self.unpack()
+        except OutOfData:
+            raise UnpackValueError("Data is not enough")
+        if self.buf_head < self.buf_tail:
+            raise ExtraData(result, self.buf[self.buf_head:])
+        return result
+
     def skip(self, object write_bytes=None):
         """
         read and ignore one object, returning None
diff --git a/test/test_sequnpack.py b/test/test_sequnpack.py
@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 # coding: utf-8
 
+import py
 import six
 from msgpack import Unpacker, BufferFull
-from msgpack.exceptions import OutOfData
+from msgpack.exceptions import OutOfData, ExtraData, UnpackValueError
 from pytest import raises
 
 
@@ -85,3 +86,15 @@ def test_readbytes():
     assert unpacker.unpack() == ord(b'a')
     assert unpacker.unpack() == ord(b'r')
 
+def test_unpack_one():
+    unpacker = Unpacker()
+    unpacker.feed('\xda\x00\x03abc')
+    assert unpacker.unpack_one() == 'abc'
+    #
+    unpacker = Unpacker()
+    unpacker.feed('\xda\x00\x03abcd')
+    py.test.raises(ExtraData, "unpacker.unpack_one()")
+    #
+    unpacker = Unpacker()
+    unpacker.feed('\xda\x00\x03ab')
+    py.test.raises(UnpackValueError, "unpacker.unpack_one()")
