importpkg: use yaml dumper directly

Instead of carefully crafting an iterator to pass to yaml.safe_dump_all,
we simply take control on our own and call represent on a yaml dumper
object where needed.

1 files changed, 13 insertions, 16 deletions

diff --git a/importpkg.py b/importpkg.py
index 2d372e4..2cdefc4 100755
--- a/importpkg.py
+++ b/importpkg.py
@@ -6,7 +6,6 @@ document contains package metadata. Then a document is emitted for each file.
 And finally a document consisting of the string "commit" is emitted."""
 
 import hashlib
-import itertools
 import optparse
 import sys
 import tarfile
@@ -55,7 +54,7 @@ def decompress_tar(filelike, extension):
     return tarfile.open(fileobj=filelike,
                         mode="r|" + extension[1:].decode("ascii"))
 
-def process_package(filelike, hash_functions):
+def process_package(filelike, hash_functions, callback):
     af = ArReader(filelike)
     af.read_magic()
     state = "start"
@@ -75,7 +74,7 @@ def process_package(filelike, hash_functions):
                 if state != "control":
                     raise ValueError("duplicate control file")
                 state = "control_file"
-                yield process_control(tf.extractfile(elem).read())
+                callback(process_control(tf.extractfile(elem).read()))
                 break
             continue
         elif name.startswith(b"data.tar"):
@@ -89,17 +88,9 @@ def process_package(filelike, hash_functions):
                 except UnicodeDecodeError:
                     print("warning: skipping filename with encoding error")
                     continue # skip files with non-utf8 encoding for now
-                yield dict(name=name, size=size, hashes=hashes)
+                callback(dict(name=name, size=size, hashes=hashes))
             break
 
-def hashed_stream_check(hstream, hashvalue):
-    if False: # pylint: disable=using-constant-test
-        yield # defer checking until being iterated
-    while hstream.read(4096):
-        pass
-    if hstream.hexdigest() != hashvalue:
-        raise ValueError("hash sum mismatch")
-
 def main():
     parser = optparse.OptionParser()
     parser.add_option("-H", "--hash", action="store",
@@ -110,12 +101,18 @@ def main():
         stdin = sys.stdin.buffer
     except AttributeError: # python2
         stdin = sys.stdin
-    iters = [("commit",)]
+    dumper = yaml.SafeDumper(sys.stdout)
+    dumper.open()
     if options.hash:
         stdin = HashedStream(stdin, hashlib.sha256())
-        iters.insert(0, hashed_stream_check(stdin, options.hash))
-    iters.insert(0, process_package(stdin, hash_functions))
-    yaml.safe_dump_all(itertools.chain(*iters), sys.stdout)
+    process_package(stdin, hash_functions, dumper.represent)
+    if options.hash:
+        while stdin.read(4096):
+            pass
+        if stdin.hexdigest() != options.hash:
+            raise ValueError("hash sum mismatch")
+    dumper.represent("commit")
+    dumper.close()
 
 if __name__ == "__main__":
     main()