def _validate_file(p):
if not path.isfile(p):
return False
ftest = root.TFile(p)
if bool(ftest) and not (ftest.IsZombie()):
logger.info(_sname + '._validate_file', '%s is a good file' % p)
return True
return False
def drop_branches(to_drop=None, to_keep=None):
if not to_drop and not to_keep:
return 0
if to_drop and to_keep:
logger.error(_sname + '.drop_branches',
'Can only provide to_drop OR to_keep')
return 0
f = root.TFile('output.root', 'UPDATE')
t = f.FindObjectAny('events')
n_entries = t.GetEntriesFast() # to check the file wasn't corrupted
if to_drop:
if type(to_drop) == str:
t.SetBranchStatus(to_drop, False)
else:
for b in to_drop:
t.SetBranchStatus(b, False)
elif to_keep:
t.SetBranchStatus('*', False)
if type(to_keep) == str:
t.SetBranchStatus(to_keep, True)
else:
for b in to_keep:
t.SetBranchStatus(b, True)
t_clone = t.CloneTree()
f.WriteTObject(t_clone, 'events', 'overwrite')
f.Close()
# check that the write went okay
f = root.TFile('output.root')
if f.IsZombie():
logger.error(_sname + '.drop_branches',
'Corrupted file trying to drop ' + str(to_drop))
return 1
t_clone = f.FindObjectAny('events')
if (n_entries == t_clone.GetEntriesFast()):
return 0
else:
logger.error(_sname + '.drop_branches',
'Corrupted tree trying to drop ' + str(to_drop))
return 2
def save(self, fpath, treename, opts='RECREATE'):
f = root.TFile(fpath, opts)
array_as_tree(self.data, treename, f)
f.Close()