simdata: load files lazily

1 files changed, 46 insertions, 3 deletions

diff --git a/simdata.py b/simdata.py
index 37bff6f..c56bd0e 100644
--- a/simdata.py
+++ b/simdata.py
@@ -3,6 +3,7 @@ This module provides a wrapper around a Cactus simulation output, i.e. a directo
 """
 
 import os
+import collections
 
 from . import datafile
 
@@ -86,6 +87,48 @@ class _RefinementLevel(object):
     def slice(self, iteration = -1, time = -1):
         return _RefinedSlice(self._sd, self, time, iteration)
 
+class _DataDir(collections.MutableMapping):
+    _files = None
+    _paths = None
+
+    def __init__(self):
+        self._files = {}
+        self._paths = {}
+
+    def __setitem__(self, key, val):
+        if key in self._files:
+            raise ValueError
+
+        self._files[key] = None
+        self._paths[key] = val
+
+    def __delitem__(self, key):
+        if self._files[key]:
+            self._files[key].close()
+        del self._files[key]
+        del self._paths[key]
+
+    def __getitem__(self, key):
+        if not key in self._paths:
+            raise KeyError
+        elif self._files[key] is None:
+            self._files[key] = datafile.DataFile(self._paths[key])
+
+        return self._files[key]
+
+    def __len__(self):
+        return len(self._paths)
+
+    def __iter__(self):
+        for it in self._paths:
+            yield it
+
+    def close(self):
+        for f in self._files.values():
+            if f is not None:
+                f.close()
+        self._files.clear()
+        self._paths.clear()
 
 class SimulationData(object):
     "The full path to the simulation directory"
@@ -101,14 +144,14 @@ class SimulationData(object):
 
     def __init__(self, dirname):
         self.dirname = os.path.abspath(dirname)
-        self.df = {}
+        self.df = _DataDir()
 
         # open all the hdf5 files in the dir
-        for f in os.listdir(dirname):
+        for f in os.listdir(self.dirname):
             if (not f.endswith('.h5') or f.startswith('checkpoint') or f.endswith('.d.h5')):
                 continue
 
-            self.df[f[:-3]] = datafile.DataFile('%s/%s' % (dirname, f))
+            self.df[f[:-3]] = '%s/%s' % (self.dirname, f)
         if len(self.df) == 0:
             raise ValueError('No HDF5 data files in the directory.')