Commit 5739edf5 authored by Merlijn Wajer's avatar Merlijn Wajer
Browse files

imagestack: add iterate_imagestack

parent a2b63f31
......@@ -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
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment