imagestack: add iterate_imagestack

derivermodule/imagestack.py

@@ -2,6 +2,9 @@ from pathlib import Path
 from subprocess import check_call
 from time import time
 
+import zipfile
+import tarfile
+
 from .logger import logger
 
 def get_imagestack_info(task_info):
@@ -76,6 +79,64 @@ def unpack_and_validate_imagestack(imagestack_path, imagestack_info, dst):
 
     return str(img_dir), image_count
 
-# TODO: function to unpack an image stack
+def iterate_imagestack(imagestack_path, imagestack_info, sort=True):
+    """Unpack and validate an imagestack
+
+    An imagestack is valid if it contains at least one directory that contains
+    at least one image of the expected image type.
+
+    Args:
+
+    * imagestack_path (``str``): The imagestack archive path
+    * imagestack_info (``dict``)::
+
+        >>> {'archive_type': ..., 'image_type': ...}
+
+    * sort (``bool``): optional, sort the files by name (sort=False should be
+                       faster if you don't care about the order)
+
+    Returns:
+
+    * Iterator of (filename, file handle)
+    """
+    if imagestack_info['archive_type'] == 'tar':
+        tf = tarfile.open(imagestack_path)
+
+        members = tf.getmembers()
+        if sort:
+            members = sorted(members, key=lambda x: x.name)
+
+        for idx, img in enumerate(members):
+            if not img.isfile():
+                continue
+
+            ci = tf.extractfile(img)
+
+            yield img.name, ci
+
+            ci.close()
+
+        tf.close()
+
+    elif imagestack_info['archive_type'] == 'zip':
+        zf = zipfile.ZipFile(imagestack_path)
+        namelist = zf.namelist()
+
+        if sort:
+            namelist = sorted(namelist)
+
+        for idx, img in enumerate(namelist):
+            info = zf.getinfo(img)
+            if info.is_dir():
+                continue
+            ci = zf.open(img)
+
+            yield info.filename, ci
+
+            ci.close()
+
+        zf.close()
+    else:
+        raise ValueError('Cannot extract archive_type %s' % imagestack_info['archive_type'])
 
-# TODO: function to iterate over images in an image stack
+    return