def _rd_delitem_(rdir, name, cycle=';*'):
"""Delete an object from ROOT file/directory
>>> del rfile['h1']
>>> del rfile['mydir/h1']
"""
##
if not rdir:
raise IOError("TDirectory is invalid")
elif not rdir.IsWritable():
raise IOError("TDirectory '%s' is not writable" % rdir.GetPath())
##
with ROOTCWD():
##
dirname, sep, fname = name.partition('/')
##
while sep:
rdir.cd()
subdir = rdir.GetDirectory(dirname)
if not subdir:
raise KeyError("TDirectory, can't delete %s" % name)
rdir = subdir
dirname, sep, fname = fname.partition('/')
rdir.cd()
if not rdir.FindKey(dirname):
raise KeyError("TDirectory, can't delete %s" % dirname)
return rdir.Delete(dirname + cycle)
def _rd_iteritems_(rdir,
fun=lambda k, t, o: True,
recursive=True,
no_dir=True):
"""Iterate over the content of ROOT directory/file:
>>> for key,obj in rfile.iteritems() : print key , obj
>>> for key,hist in rfile.iteritems( ROOT.TH1 ) : print key , hist
>>> for key,obj in rfile.iteritems( lambda name,tkey,obj : name[0]=='M' ) : print key,obj
"""
##
if isinstance(fun, type) and issubclass(fun, (ROOT.TObject, cpp.TObject)):
tobj = fun
fun = lambda k, t, o: isinstance(o, tobj)
##
with ROOTCWD():
##
rdir.cd()
_lst = rdir.GetListOfKeys()
for i in _lst:
inam = i.GetName()
idir = rdir.GetDirectory(inam)
if not idir or not no_dir:
obj = rdir.Get(inam)
if fun(inam, i, obj): yield inam, obj
if recursive and idir:
for k, o in _rd_iteritems_(idir, fun, recursive, no_dir):
yield k, o
def _rd_getitem_(rdir, name):
"""Get the object from ROOT-file/directory
>>> t = f['A/tup']
>>> h = f['histo']
"""
##
if not rdir: raise IOError("TDirectory is invalid")
##
with ROOTCWD():
##
dirname, sep, fname = name.partition('/')
##
while sep:
rdir.cd()
subdir = rdir.GetDirectory(dirname)
if not subdir:
raise KeyError("TDirectory[]: unknown directory name '%s%s'" %
(rdir.GetPath(), dirname))
rdir = subdir
dirname, sep, fname = fname.partition('/')
rdir.cd()
obj = rdir.Get(dirname)
if not obj:
raise KeyError("TDirectory[]: unknown object '%s%s'" %
(rdir.GetPath(), dirname))
return obj
def _rd_setitem_(rdir, name, tobj):
"""Write the object to ROOT-file/directory
>>> rfile['myhisto'] = h1
>>> rfile['MyDir/MyHisto'] = h2
"""
##
if not rdir:
raise IOError("TDirectory is invalid")
elif not rdir.IsWritable():
raise IOError("TDirectory '%s' is not writable" % rdir.GetPath())
##
with ROOTCWD():
##
dirname, sep, fname = name.partition('/')
##
while sep:
rdir.cd()
subdir = rdir.GetDirectory(dirname)
if not subdir:
rdir.cd()
subdir = rdir.mkdir(dirname, rdir.GetPath() + '/' + dirname)
rdir.cd()
subdir = rdir.GetDirectory(dirname)
if not subdir:
raise KeyError("TDirectory[]: can't create '%s%s'" %
(rdir.GetPath(), dirname))
rdir = subdir
dirname, sep, fname = fname.partition('/')
rdir.cd()
return rdir.WriteTObject(tobj, dirname, 'WriteDelete')
def _rd_iter_(rdir):
"""Iterate (recursive) over the content in ROOT file/directory
>>> for obj in rfile : print obj
"""
##
with ROOTCWD():
##
rdir.cd()
for obj in _rd_itervalues_(rdir, recursive=True):
yield obj
def _rf_new_close_(rfile, options=''):
"""
Open/create ROOT-file without making it a current working directory
>>> print ROOT.gROOT.CurrentDirectory()
>>> f = ROOT.TFile.Open('test_file.root','recreate')
>>> print ROOT.gROOT.CurrentDirectory()
"""
if rfile and rfile.IsOpen():
with ROOTCWD():
##rfile .cd()
logger.debug("Close ROOT file %s" % rfile.GetName())
rfile._old_close_(options)
def _rf_new_init_(rfile, fname, mode='', *args):
"""Open/create ROOT-file without making it a current working directory
>>> print ROOT.gROOT.CurrentDirectory()
>>> f = ROOT.TFile('test_file.root','recreate')
>>> print ROOT.gROOT.CurrentDirectory()
"""
with ROOTCWD():
logger.debug("Open ROOT file %s/'%s'" % (fname, mode))
result = rfile._old_init_(fname, open_mode(mode), *args)
if not rfile or not rfile.IsOpen():
logger.error("Can't open ROOT file %s/'%s'" % (fname, mode))
return result
def _rd_itervalues_(rdir,
fun=lambda k, t, o: True,
recursive=True,
no_dir=True):
"""Iterate over the content of ROOT directory/file:
>>> for obj in rfile.itervalues () : print obj
>>> for hist in rfile.itervalues ( ROOT.TH1 ) : print hist
>>> for obj in rfile.itervalues ( lambda k,t,o : k[0]=='M' ) : print obj
"""
##
with ROOTCWD():
##
rdir.cd()
for key, obj in _rd_iteritems_(rdir, fun, recursive, no_dir):
yield obj
def _rd_iterkeys_(rdir, typ=None, recursive=True, no_dir=True):
"""Iterate over the keys in ROOT file/directory
>>> for key,obj in rfile.iteritems() : print key , obj
>>> for key,hist in rfile.iteritems( ROOT.TH1 ) : print key , hist
"""
##
with ROOTCWD():
##
rdir.cd()
_lst = rdir.GetListOfKeys()
for i in _lst:
inam = i.GetName()
idir = rdir.GetDirectory(inam)
if not idir or not no_dir:
if typ is None or isinstance(rdir.Get(inam), typ): yield inam
if recursive and idir:
for k in _rd_iterkeys_(idir, typ, recursive, no_dir):
yield k
def _rd_contains_(rdir, name):
"""Check the existence of key in the root-file/directory
>>> if not 'myhisto' in rfile : break
>>> if 'MyDir/MyHisto' in rfile : print 'OK!'
"""
##
if not rdir: return False
##
with ROOTCWD():
##
dirname, sep, fname = name.partition('/')
##
while sep:
rdir.cd()
subdir = rdir.GetDirectory(dirname)
if not subdir: return False ## RETURN
rdir = subdir
dirname, sep, fname = fname.partition('/')
rdir.cd()
return rdir.FindKey(dirname) ## RETURN
def _rd_keys_(rdir, recursive=True, no_dir=True):
"""Get all keys from ROOT file/directory
>>> keys = rfile.keys()
"""
_res = []
if not rdir: return _res
##
with ROOTCWD():
##
rdir.cd()
_lst = rdir.GetListOfKeys()
for i in _lst:
inam = i.GetName()
idir = rdir.GetDirectory(inam)
if not idir or not no_dir: _res.append(inam)
if recursive and idir:
ikeys = _rd_keys_(idir, recursive, no_dir)
for k in ikeys:
_res.append(inam + '/' + k)
return _res
def _rd_getattr_(rdir, name):
"""
Get the object from ROOT-file
>>> t = f.A.tup
>>> h = f.myhisto
"""
##
if not rdir: return None
##
with ROOTCWD():
##
dirname, sep, fname = name.partition('/')
if sep:
raise AttributeError(
'TDirectory[]: invalid attribute %s.%s' % rdir.GetPath(),
dirname)
rdir.cd()
obj = rdir.Get(dirname)
if not obj:
raise AttributeError('TDirectory[]: unknown attribute %s.%s' %
(rdir.GetPath(), dirname))
return obj