Add a thing to make using BufferCollator as some sort of queue easy and ...

docs/api-reference.rst

@@ -187,9 +187,16 @@ API Reference
 
         Adds the contents of a view to the collator.
 
-    .. method:: collapse()
+    .. method:: collapse(max_bytes=None)
+
+        :param int max_bytes:
 
         Collapses the contents of the collator into a single
         :class:`BufferView`. Also resets the internal state of the collator, so
         if you call it twice successively, the second call will return an empty
         :class:`BufferView`.
+
+        If ``max_bytes`` is provided, only the first ``max_bytes`` of the
+        contents are copied into the new :class:`Buffer`, leaving the remainder
+        in the collator. If ``max_bytes`` is larger than then length of the
+        collator, it is capped to ``len(collator)``.

test_zero_buffer.py

@@ -359,7 +359,7 @@ def test_single_item(self, buf):
         view = buf.view()
         collator = BufferCollator()
         collator.append(view)
-        assert collator.collapse() is view
+        assert collator.collapse()._data == view._data
 
     def test_collapse_clears(self, buf):
         buf.add_bytes(b"abc")
@@ -377,3 +377,33 @@ def test_len(self, buf):
         collator.append(view)
         collator.append(view)
         assert len(collator) == 6
+
+    def test_max_bytes(self, buf):
+        buf.add_bytes(b"abc")
+        view = buf.view()
+        collator = BufferCollator()
+        collator.append(view)
+        collator.append(view)
+        view = collator.collapse(max_bytes=4)
+        assert view == b"abca"
+        view = collator.collapse()
+        assert view == b"bc"
+
+    def test_max_bytes_partial_first_item(self, buf):
+        buf.add_bytes(b"abc")
+        view = buf.view()
+        collator = BufferCollator()
+        collator.append(view)
+        collator.append(view)
+        view = collator.collapse(max_bytes=2)
+        assert view == b"ab"
+        next_view = collator.collapse(max_bytes=4)
+        assert next_view == b"cabc"
+
+    def test_max_bytes_larger_than_length(self, buf):
+        buf.add_bytes(b"abc")
+        view = buf.view()
+        collator = BufferCollator()
+        collator.append(view)
+        view = collator.collapse(max_bytes=100)
+        assert view == b"abc"

zero_buffer.py

@@ -14,7 +14,6 @@
 int memcmp(const void *, const void *, size_t);
 void *memchr(const void *, int, size_t);
 void *Zero_memrchr(const void *, int, size_t);
-void *memcpy(void *, const void *, size_t);
 """)
 _lib = _ffi.verify("""
 #include <string.h>
@@ -453,16 +452,30 @@ def append(self, view):
             self._views.append(view)
         self._total_length += len(view)
 
-    def collapse(self):
+    def collapse(self, max_bytes=None):
+        if max_bytes is None:
+            max_bytes = self._total_length
+        else:
+            max_bytes = min(self._total_length, max_bytes)
+
         if len(self._views) == 1:
-            result = self._views[0]
+            result = self._views[0][:max_bytes]
+            del self._views[0]
         else:
-            data = _ffi.new("uint8_t[]", self._total_length)
+            data = _ffi.new("uint8_t[]", max_bytes)
             pos = 0
-            for view in self._views:
-                _lib.memcpy(data + pos, view._data, len(view))
+            for i, view in enumerate(self._views):
+                if len(view) > max_bytes - pos:
+                    data[pos:max_bytes] = view._data[0:max_bytes - pos]
+                    if i >= 1:
+                        del self._views[:i - 1]
+                    self._views[0] = self._views[0][max_bytes - pos:]
+                    break
+                else:
+                    data[pos:pos + len(view)] = view._data[0:len(view)]
                 pos += len(view)
-            result = Buffer(data, self._total_length).view()
-        del self._views[:]
-        self._total_length = 0
+            else:
+                del self._views[:]
+            result = Buffer(data, max_bytes).view()
+        self._total_length -= max_bytes
         return result