DecompressedStream: implement readline

Iteration over file-like is required by deb822.Packages.iter_paragraphs.

1 files changed, 27 insertions, 10 deletions

diff --git a/dedup/compression.py b/dedup/compression.py
index 7f6dc99..ea431cc 100644
--- a/dedup/compression.py
+++ b/dedup/compression.py
@@ -108,26 +108,43 @@ class DecompressedStream(object):
         self.pos = 0
         self.closed = False
 
-    def read(self, length=None):
+    def _fill_buff_until(self, predicate):
         assert not self.closed
         data = True
         while True:
-            if length is not None and len(self.buff) >= length:
-                ret = self.buff[:length]
-                self.buff = self.buff[length:]
-                break
-            elif not data: # read EOF in last iteration
-                ret = self.buff
-                self.buff = b""
-                break
+            if predicate(self.buff) or not data:
+                return
             data = self.fileobj.read(self.blocksize)
             if data:
                 self.buff += self.decompressor.decompress(data)
             else:
                 self.buff += self.decompressor.flush()
-        self.pos += len(ret)
+
+    def _read_from_buff(self, length):
+        ret = self.buff[:length]
+        self.buff = self.buff[length:]
+        self.pos += length
         return ret
 
+    def read(self, length=None):
+        if length is None:
+            self._fill_buff_until(lambda _: False)
+            length = len(self.buff)
+        else:
+            self._fill_buff_until(lambda b, l=length: len(b) >= l)
+        return self._read_from_buff(length)
+
+    def readline(self):
+        self._fill_buff_until(lambda b: b'\n' in b)
+        try:
+            length = self.buff.index(b'\n') + 1
+        except ValueError:
+            length = len(self.buff)
+        return self._read_from_buff(length)
+
+    def __iter__(self):
+        return iter(self.readline, b'')
+
     def tell(self):
         assert not self.closed
         return self.pos