importpkg: refactor commit handling out of process_package*

1 files changed, 13 insertions, 16 deletions

diff --git a/importpkg.py b/importpkg.py
index 99e9850..2d372e4 100755
--- a/importpkg.py
+++ b/importpkg.py
@@ -6,6 +6,7 @@ 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
@@ -89,20 +90,15 @@ def process_package(filelike, hash_functions):
                     print("warning: skipping filename with encoding error")
                     continue # skip files with non-utf8 encoding for now
                 yield dict(name=name, size=size, hashes=hashes)
-            yield "commit"
             break
 
-def process_package_with_hash(filelike, hash_functions, sha256hash):
-    hstream = HashedStream(filelike, hashlib.sha256())
-    for elem in process_package(hstream, hash_functions):
-        if elem == "commit":
-            while hstream.read(4096):
-                pass
-            if hstream.hexdigest() != sha256hash:
-                raise ValueError("hash sum mismatch")
-            yield elem
-            break
-        yield elem
+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()
@@ -114,11 +110,12 @@ def main():
         stdin = sys.stdin.buffer
     except AttributeError: # python2
         stdin = sys.stdin
+    iters = [("commit",)]
     if options.hash:
-        gen = process_package_with_hash(stdin, hash_functions, options.hash)
-    else:
-        gen = process_package(stdin, hash_functions)
-    yaml.safe_dump_all(gen, sys.stdout)
+        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)
 
 if __name__ == "__main__":
     main()