Rewrite the HDF wrapping code.

1 files changed, 37 insertions, 18 deletions

diff --git a/cactus_utils.py b/cactus_utils.py
index a68767c..867b175 100644
--- a/cactus_utils.py
+++ b/cactus_utils.py
@@ -976,34 +976,53 @@ def l2norm(sd, it, gf, maxrl = None, minrl = None):
     return sum(results)
 
 
-class Dataset:
+class HDFObjWrapper:
+    ### private ###
+    _file = None
+    _obj  = None
+
+    def __init__(self, obj, file):
+        self._file = file
+        self._obj  = obj
+
+    def _wrap(self, item):
+        if isinstance(item, h5py.Group) or isinstance(item, h5py.Dataset):
+            return HDFObjWrapper(item, self._file)
+        return item
+
+    def __getitem__(self, key):
+        return self._wrap(self._obj[key])
+
+    @property
+    def attrs(self):
+        return self._wrap(self._obj.attrs)
+
+    @property
+    def value(self):
+        return self._wrap(self._obj.value)
+
+
+class HDFSource(dict):
     """
-    A wrapper around the h5py Dataset that holds a reference to the underlying
-    file and closes it on destruct.
+    A wrapper around a HDF file opened in read-only mode, making sure that the
+    file is automatically closed when the instance is destroyed.
     """
     ### private ###
-    # the wrapped dataset
-    _dataset = None
-    # the underlying File object
     _file = None
 
-    def __init__(self, dataset, file):
-        self.__dict__['_dataset'] = dataset
-        self.__dict__['_file']    = file
-
-    def __getattr__(self, name):
-        attr = self._dataset.__getattribute__(name)
-        if callable(attr):
-            def call_wrapper(*args, **kwargs):
-                return attr(*args, **kwargs)
-            return call_wrapper
-        else:
-            return attr
+    def __init__(self, path):
+        self._file = h5py.File(path, 'r')
 
     def __del__(self):
         if self._file is not None:
             self._file.close()
 
+    def __getitem__(self, key):
+        return HDFObjWrapper(self._file[key], self)
+
+    def __iter__(self):
+        return self._file.__iter__()
+
 
 class RefinedSlice:
     # simulation time on this slice