Rewrite generator to IOStream

koji/__init__.py

@@ -57,6 +57,7 @@ except ImportError:  # pragma: no cover
 from fnmatch import fnmatch
 
 import dateutil.parser
+import io
 import requests
 import six
 import six.moves.configparser
@@ -945,23 +946,47 @@ def get_sighdr_key(sighdr):
 
 def spliced_sig_reader(path, sighdr, bufsize=8192):
     """A generator that yields the contents of an rpm with signature spliced in"""
-    (start, size) = find_rpm_sighdr(path)
-    with open(path, 'rb') as fo:
-        # the part before the signature
-        yield fo.read(start)
+    class Stream(io.RawIOBase):
+        def __init__(self, path, sighdr):
+            self.path = path
+            self.sighdr = sighdr
+            self.buf = None
+            self.gen = self.generator()
 
-        # the spliced signature
-        yield sighdr
+        def generator(self):
+            (start, size) = find_rpm_sighdr(self.path)
+            with open(path, 'rb') as fo:
+                # the part before the signature
+                yield fo.read(start)
 
-        # skip original signature
-        fo.seek(size, 1)
+                # the spliced signature
+                yield sighdr
 
-        # the part after the signature
-        while True:
-            buf = fo.read(bufsize)
-            if not buf:
-                break
-            yield buf
+                # skip original signature
+                fo.seek(size, 1)
+
+                # the part after the signature
+                while True:
+                    buf = fo.read(bufsize)
+                    if not buf:
+                        break
+                    yield buf
+
+        def readable(self):
+            return True
+
+        def readinto(self, b):
+            try:
+                expected_buf_size = len(b)
+                data = self.buf or next(self.gen)
+                output = data[:expected_buf_size]
+                self.buf = data[expected_buf_size:]
+                b[:len(output)] = output
+                return len(output)
+            except StopIteration:
+                return 0    # indicate EOF
+
+    return io.BufferedReader(Stream(path, sighdr), buffer_size=bufsize)
 
 
 def splice_rpm_sighdr(sighdr, src, dst=None, bufsize=8192, callback=None):