Support flat files and hdf5

gistfile1.py

import os
import h5py
import numpy as np


def touch(fname, data):
    f = open(fname, 'w')
    f.close()


class SequenceWriter(object):

    def __init__(self, prefix, write_func):
        """
        *write_func* is a callable that receives a file name prefix and
        data, e.g. write_tiff().
        """

        self.prefix = prefix
        self.write_func = write_func
        self.count = 0

    def __call__(self, data):
        self.write_func(self.prefix.format(self.count), data)
        self.count += 1


class Walker(object):

    def descent(self, name):
        raise NotImplementedError

    def ascent(self):
        raise NotImplementedError

    def write_sequence(self, generator):
        raise NotImplementedError

class DirectoryWalker(Walker):

    def __init__(self, write_func, fname_fmt):
        self._root = os.getcwd()
        self._current = self._root
        self._fname_fmt = fname_fmt
        self._write_func = write_func

    def descent(self, name):
        self._current = os.path.join(self._current, name)

        if not os.path.exists(self._current):
            os.mkdir(self._current)

    def ascent(self):
        if self._current != self._root:
            self._current = os.path.dirname(self._current)

    def write_sequence(self, generator):
        path = os.path.join(self._current, self._fname_fmt)
        writer = SequenceWriter(path, self._write_func)

        for data in generator:
            writer(data)


class Hdf5Walker(Walker):

    def __init__(self, hdf5, fname):
        self._root = hdf5
        self._current = self._root
        self._fname = fname

    def descent(self, name):
        self._current = self._current.create_group(name)

    def ascent(self):
        self._current = self._current.parent

    def write_sequence(self, generator):
        # this is *very* inefficient ...
        data = np.array(list(generator))
        dset = self._current.create_dataset(self._fname, data.shape, dtype=data.dtype)

        # this is specific to sequences of images and not very general
        dset[:,:,:] = data[:,:,:]


def save_data(walker, data):
    walker.descent('foo')
    walker.descent('bar')
    walker.write_sequence(data)
    walker.ascent()
    walker.descent('baz')


if __name__ == '__main__':
    dw = DirectoryWalker(touch, 'qux-{:>05}.tif')
    hw = Hdf5Walker(h5py.File('foo.hdf5', 'w'), 'qux')

    data = [np.random.random((16, 16)) for i in range(4)]
    save_data(dw, data)
    save_data(hw, data)