h5config/storage/hdf5.py

   1 # Copyright (C) 2011 W. Trevor King <wking@drexel.edu>
   2 #
   3 # This file is part of h5config.
   4 #
   5 # h5config is free software; you can redistribute it and/or modify it
   6 # under the terms of the GNU General Public License as published by the
   7 # Free Software Foundation, either version 3 of the License, or (at your
   8 # option) any later version.
   9 #
  10 # h5config is distributed in the hope that it will be useful, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of
  12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  13 # General Public License for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License
  16 # along with h5config.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 """HDF5 backend implementation
  19 """
  20
  21 import types as _types
  22
  23 import h5py as _h5py
  24
  25 from .. import LOG as _LOG
  26 from .. import config as _config
  27 from . import FileStorage as _FileStorage
  28
  29
  30 def pprint_HDF5(*args, **kwargs):
  31     print pformat_HDF5(*args, **kwargs)
  32
  33 def pformat_HDF5(filename, group='/'):
  34     with _h5py.File(filename, 'r') as f:
  35         cwg = f[group]
  36         ret = '\n'.join(_pformat_hdf5(cwg))
  37     return ret
  38
  39 def _pformat_hdf5(cwg, depth=0):
  40     lines = []
  41     lines.append('  '*depth + cwg.name)
  42     depth += 1
  43     for key,value in cwg.iteritems():
  44         if isinstance(value, _h5py.Group):
  45             lines.extend(_pformat_hdf5(value, depth))
  46         elif isinstance(value, _h5py.Dataset):
  47             lines.append('  '*depth + str(value))
  48             lines.append('  '*(depth+1) + str(value[...]))
  49         else:
  50             lines.append('  '*depth + str(value))
  51     return lines
  52
  53 def h5_create_group(cwg, path, force=False):
  54     "Create the group where the settings are stored (if necessary)."
  55     if path == '/':
  56         return cwg
  57     gpath = ['']
  58     for group in path.strip('/').split('/'):
  59         gpath.append(group)
  60         if group not in cwg.keys():
  61             _LOG.debug('creating group {} in {}'.format(
  62                     '/'.join(gpath), cwg.file))
  63             cwg.create_group(group)
  64         _cwg = cwg[group]
  65         if isinstance(_cwg, _h5py.Dataset):
  66             if force:
  67                 _LOG.info('overwrite {} in {} ({}) with a group'.format(
  68                         '/'.join(gpath), _cwg.file, _cwg))
  69                 del cwg[group]
  70                 _cwg = cwg.create_group(group)
  71             else:
  72                 raise ValueError(_cwg)
  73         cwg = _cwg
  74     return cwg
  75
  76
  77 class HDF5_Storage (_FileStorage):
  78     """Back a `Config` class with an HDF5 file.
  79
  80     The `.save` and `.load` methods have an optional `group` argument
  81     that allows you to save and load settings from an externally
  82     opened HDF5 file.  This can make it easier to stash several
  83     related `Config` classes in a single file.  For example
  84
  85     >>> import os
  86     >>> import tempfile
  87     >>> from ..test import TestConfig
  88     >>> fd,filename = tempfile.mkstemp(
  89     ...     suffix='.'+HDF5_Storage.extension, prefix='pypiezo-')
  90     >>> os.close(fd)
  91
  92     >>> f = _h5py.File(filename, 'a')
  93     >>> c = TestConfig(storage=HDF5_Storage(
  94     ...     filename='untouched_file.h5', group='/untouched/group'))
  95     >>> c['alive'] = True
  96     >>> group = f.create_group('base')
  97     >>> c.save(group=group)
  98     >>> pprint_HDF5(filename)  # doctest: +REPORT_UDIFF
  99     /
 100       /base
 101         <HDF5 dataset "age": shape (), type "<f8">
 102           1.3
 103         <HDF5 dataset "alive": shape (), type "|b1">
 104           True
 105         <HDF5 dataset "bids": shape (3,), type "<f8">
 106           [ 5.4  3.2  1. ]
 107         <HDF5 dataset "children": shape (), type "|S1">
 108     <BLANKLINE>
 109         <HDF5 dataset "daisies": shape (), type "<i4">
 110           13
 111         <HDF5 dataset "name": shape (), type "|S1">
 112     <BLANKLINE>
 113         <HDF5 dataset "species": shape (), type "|S14">
 114           Norwegian Blue
 115         <HDF5 dataset "spouse": shape (), type "|S1">
 116     <BLANKLINE>
 117     >>> c.clear()
 118     >>> c['alive']
 119     False
 120     >>> c.load(group=group)
 121     >>> c['alive']
 122     True
 123
 124     >>> f.close()
 125     >>> os.remove(filename)
 126     """
 127     extension = 'h5'
 128
 129     def __init__(self, group='/', **kwargs):
 130         super(HDF5_Storage, self).__init__(**kwargs)
 131         assert group.startswith('/'), group
 132         if not group.endswith('/'):
 133             group += '/'
 134         self.group = group
 135         self._file_checked = False
 136
 137     def _check_file(self):
 138         if self._file_checked:
 139             return
 140         self._setup_file()
 141         self._file_checked = True
 142
 143     def _setup_file(self):
 144         with _h5py.File(self._filename, 'a') as f:
 145             cwg = f  # current working group
 146             h5_create_group(cwg, self.group)
 147
 148     def _load(self, config, group=None):
 149         f = None
 150         try:
 151             if group is None:
 152                 self._check_file()
 153                 f = _h5py.File(self._filename, 'r')
 154                 group = f[self.group]
 155             for s in config.settings:
 156                 if s.name not in group.keys():
 157                     continue
 158                 elif isinstance(s, (_config.BooleanSetting,
 159                                     _config.NumericSetting,
 160                                     _config.FloatListSetting)):
 161                     v = group[s.name][...]
 162                     if isinstance(v, _types.StringTypes):
 163                         # convert back from None, etc.
 164                         v = s.convert_from_text(v)
 165                     elif isinstance(s, _config.FloatListSetting):
 166                         v = list(v)  # convert from numpy array
 167                     config[s.name] = v
 168                 elif isinstance(s, _config.ConfigListSetting):
 169                     try:
 170                         cwg = h5_create_group(group, s.name)
 171                     except ValueError:
 172                         pass
 173                     else:
 174                         value = []
 175                         for i in sorted(int(x) for x in cwg.keys()):
 176                             instance = s.config_class()
 177                             try:
 178                                 _cwg = h5_create_group(cwg, str(i))
 179                             except ValueError:
 180                                 pass
 181                             else:
 182                                 self._load(config=instance, group=_cwg)
 183                                 value.append(instance)
 184                         config[s.name] = value
 185                 elif isinstance(s, _config.ConfigSetting):
 186                     try:
 187                         cwg = h5_create_group(group, s.name)
 188                     except ValueError:
 189                         pass
 190                     else:
 191                         if not config[s.name]:
 192                             config[s.name] = s.config_class()
 193                         self._load(config=config[s.name], group=cwg)
 194                 else:
 195                     config[s.name] = s.convert_from_text(group[s.name][...])
 196         finally:
 197             if f:
 198                 f.close()
 199
 200     def _save(self, config, group=None):
 201         f = None
 202         try:
 203             if group is None:
 204                 self._check_file()
 205                 f = _h5py.File(self._filename, 'a')
 206                 group = f[self.group]
 207             for s in config.settings:
 208                 value = None
 209                 if isinstance(s, (_config.BooleanSetting,
 210                                   _config.NumericSetting,
 211                                   _config.FloatListSetting)):
 212                     value = config[s.name]
 213                     if value in [None, []]:
 214                         value = s.convert_to_text(value)
 215                 elif isinstance(s, _config.ConfigListSetting):
 216                     configs = config[s.name]
 217                     if configs:
 218                         cwg = h5_create_group(group, s.name, force=True)
 219                         for i,cfg in enumerate(configs):
 220                             _cwg = h5_create_group(cwg, str(i), force=True)
 221                             self._save(config=cfg, group=_cwg)
 222                         continue
 223                 elif isinstance(s, _config.ConfigSetting):
 224                     cfg = config[s.name]
 225                     if cfg:
 226                         cwg = h5_create_group(group, s.name, force=True)
 227                         self._save(config=cfg, group=cwg)
 228                         continue
 229                 if value is None:  # not set yet, or invalid
 230                     value = s.convert_to_text(config[s.name])
 231                 try:
 232                     del group[s.name]
 233                 except KeyError:
 234                     pass
 235                 group[s.name] = value
 236         finally:
 237             if f:
 238                 f.close()