ioperf/hdf5/test_modification.py

import h5py
import logging
import argparse
import time
import contextlib



@contextlib.contextmanager
def elapsed():
    start = time.time()
    yield
    end = time.time()
    logging.info("{:.3f} seconds elapsed".format(end - start))


def write(name, size):
    with elapsed():
        f = h5py.File(name, 'w')
        dset = f.create_dataset('dset', (size, ), dtype='u1')
        dset.attrs.create('name', 'some name')
        dset.attrs.create('number', 42)
        dset[:] = 123
        f.flush()
        f.close()
        logging.info("Wrote {} bytes to {}".format(size, name))


def modify(name):
    with elapsed():
        f = h5py.File(name, 'r+')
        dset = f['dset']
        dset.attrs.modify('name', 'foo')
        dset.attrs.modify('number', 23)
        f.flush()
        f.close()
        logging.info("Modified {}".format(name))


def main():
    logging.basicConfig(level=logging.DEBUG)

    parser = argparse.ArgumentParser()
    parser.add_argument('--write', action='store_true')
    parser.add_argument('--modify', action='store_true')
    args = parser.parse_args()

    files = [('test.1K.h5', 1024),
             ('test.1M.h5', 1024 * 1024),
             ('test.10M.h5', 10 * 1024 * 1024),
             ('test.100M.h5', 100 * 1024 * 1024),
             ('test.1G.h5', 1024 * 1024 * 1024)]

    if args.write:
        for name, size in files:
            write(name, size)

    if args.modify:
        for name, _ in files:
            modify(name)


if __name__ == '__main__':
    main()