importpkg: reduce copy&paste

1 files changed, 15 insertions, 33 deletions

diff --git a/importpkg.py b/importpkg.py
index 247965f..06d9da4 100755
--- a/importpkg.py
+++ b/importpkg.py
@@ -44,6 +44,15 @@ def gifhash():
     hashobj.name = "gif_sha512"
     return hashobj
 
+def decompress_tar(filelike, extension):
+    if extension in (".lzma", ".xz"):
+        filelike = DecompressedStream(filelike, lzma.LZMADecompressor())
+        extension = ""
+    if extension not in ("", ".gz", ".bz2"):
+        raise ValueError("unknown compression format with extension %r" %
+                         extension)
+    return tarfile.open(fileobj=filelike, mode="r|" + extension[1:])
+
 def process_package(filelike, hash_functions):
     af = ArReader(filelike)
     af.read_magic()
@@ -53,39 +62,11 @@ def process_package(filelike, hash_functions):
             name = af.read_entry()
         except EOFError:
             raise ValueError("data.tar not found")
-        if name == "control.tar.gz":
-            new_state = "control"
-            tf = tarfile.open(fileobj=af, mode="r|gz")
-        elif name == "control.tar.xz":
-            new_state = "control"
-            zf = DecompressedStream(af, lzma.LZMADecompressor())
-            tf = tarfile.open(fileobj=zf, mode="r|")
-        elif name == "control.tar":
-            new_state = "control"
-            tf = tarfile.open(fileobj=af, mode="r|")
-        elif name == "data.tar.gz":
-            new_state = "data"
-            tf = tarfile.open(fileobj=af, mode="r|gz")
-        elif name == "data.tar.bz2":
-            new_state = "data"
-            tf = tarfile.open(fileobj=af, mode="r|bz2")
-        elif name == "data.tar.lzma":
-            new_state = "data"
-            zf = DecompressedStream(af, lzma.LZMADecompressor())
-            tf = tarfile.open(fileobj=zf, mode="r|")
-        elif name == "data.tar.xz":
-            new_state = "data"
-            zf = DecompressedStream(af, lzma.LZMADecompressor())
-            tf = tarfile.open(fileobj=zf, mode="r|")
-        elif name == "data.tar":
-            new_state = "data"
-            tf = tarfile.open(fileobj=af, mode="r|")
-        else:
-            continue
-        if new_state == "control":
+        if name.startswith("control.tar"):
             if state != "start":
                 raise ValueError("unexpected control.tar")
-            state = new_state
+            state = "control"
+            tf = decompress_tar(af, name[11:])
             for elem in tf:
                 if elem.name != "./control":
                     continue
@@ -95,10 +76,11 @@ def process_package(filelike, hash_functions):
                 yield process_control(tf.extractfile(elem).read())
                 break
             continue
-        elif new_state == "data":
+        elif name.startswith("data.tar"):
             if state != "control_file":
                 raise ValueError("missing control file")
-            state = new_state
+            state = "data"
+            tf = decompress_tar(af, name[8:])
             for name, size, hashes in get_tar_hashes(tf, hash_functions):
                 try:
                     name = name.decode("utf8")