00001 import h5py
00002 import numpy
00003 import logging
00004
00005 class ConstantsStore(object):
00006 def __init__(self,obj,file):
00007 self.f = h5py.File(file,'w')
00008 self.cwd = ''
00009 for k in obj.keys():
00010 subobj = obj[k]
00011 self.dispatch(subobj,str(k))
00012 self.f.close()
00013 def pushdir(self,dir):
00014 '''move down a level and keep track of what hdf directory level we are in'''
00015
00016 self.cwd += '/'+dir
00017
00018 def popdir(self):
00019 '''move up a level and keep track of what hdf directory level we are in'''
00020 self.cwd = self.cwd[:self.cwd.rfind('/')]
00021 def typeok(self,obj,name):
00022 '''check if we support serializing this type to hdf'''
00023 allowed = [dict,int,float,str,numpy.ndarray]
00024 return type(obj) in allowed
00025 def storevalue(self,v,name):
00026 '''persist one of the supported types to the hdf file'''
00027 self.f[self.cwd+'/'+name] = v
00028 def dict(self,d,name):
00029 '''called for every dictionary level to create a new hdf group name.
00030 it then looks into the dictionary to see if other groups need to
00031 be created'''
00032 if self.cwd is '':
00033 self.f.create_group(name)
00034 self.pushdir(name)
00035 for k in d.keys():
00036 self.dispatch(d[k],str(k))
00037 self.popdir()
00038 def dispatch(self,obj,name):
00039 '''either persist a supported object, or look into a dictionary
00040 to see what objects need to be persisted'''
00041 if type(obj) is dict:
00042 self.dict(obj,name)
00043 else:
00044 if self.typeok(obj,name):
00045 self.storevalue(obj,name)
00046 else:
00047 logging.warning('Constants.py: variable "'+name+'" of type "'+type(obj).__name__+'" not supported')
00048
00049 class ConstantsLoad(object):
00050 def __init__(self,file):
00051 self.obj = {}
00052 self.f = h5py.File(file,'r')
00053 self.f.visititems(self.loadCallBack)
00054 self.f.close()
00055 def setval(self,name,obj):
00056 '''see if this hdfname has a / in it. if so, create the dictionary
00057 object. if not, set our attribute value. call ourselves
00058 recursively to see if other dictionary levels exist.'''
00059 if '/' in name:
00060 dictname=name[:name.find('/')]
00061 remainder=name[name.find('/')+1:]
00062
00063 if not dictname in obj:
00064 obj[dictname]={}
00065
00066 self.setval(remainder,obj[dictname])
00067 else:
00068 obj[name]=self.f[self.fullname].value
00069
00070 def loadCallBack(self,name,obj):
00071 '''called back by h5py routine visititems for each
00072 item (group/dataset) in the h5 file'''
00073 if isinstance(obj,h5py._hl.group.Group):
00074 return
00075 self.fullname = name
00076 self.setval(name,self.obj)
00077
00078 def load(file):
00079 '''takes a string filename, and returns a constants object.'''
00080 c = ConstantsLoad(file)
00081 return c.obj
00082
00083 def save(file,obj):
00084 '''store a constants object in an hdf5 file. the object
00085 can be a hierarchy (defined by python dictionaries) and
00086 hdf5 supported types (int, float, numpy.ndarray, string).
00087 the hierarchy can be created by having one value of
00088 a dictionary itself be a dictionary.'''
00089
00090 c = ConstantsStore(obj,file)
00091
00092 if __name__ == "__main__":
00093 ct = { 'version' : 0,
00094 'darkreferencepath':'hello',
00095 'nb':12,
00096 'subdict':{'first' : 1, 'second' : 'two','three' : 'bahahah'}
00097 }
00098 save('ConstTest.h5',ct)
00099 data = load('ConstTest.h5')
00100 print '***',data
