8 __doc__ =
"a few utils to ease the day-to-day work with ROOT"
9 __author__ =
"Sebastien Binet"
19 from functools
import cache
23 """a helper method to wrap the 'import ROOT' statement to prevent ROOT
24 from screwing up the display or loading graphics libraries when in batch
25 mode (which is the default.)
28 >>> ROOT = import_root(batch=True)
29 >>> f = ROOT.TFile.Open(...)
32 ROOT.gROOT.SetBatch(batch)
34 ROOT.PyConfig.IgnoreCommandLineOptions =
True
36 if os.environ.get(
'GLIBCXX_USE_CXX11_ABI') ==
'0':
37 cmd = ROOT.gSystem.GetMakeSharedLib()
38 if cmd.find(
'GLIBCXX_USE_CXX11_ABI') < 0:
39 cmd = cmd.replace (
'$SourceFiles',
'$SourceFiles -D_GLIBCXX_USE_CXX11_ABI=0 ')
40 ROOT.gSystem.SetMakeSharedLib(cmd)
46 """a helper method to compile a set of C++ statements (via ``src``) or
47 a C++ file (via ``fname``) via ACLiC
49 if src
is not None and fname
is not None:
50 raise ValueError(
"'src' xor 'fname' should be not None, *not* both")
52 if src
is None and fname
is None:
53 raise ValueError(
"'src' xor 'fname' should be None, *not* both")
62 _first_compile =
False
63 root_compile (
'extern "C" { void __gmon_start__(){}; }',
None,
True)
64 return _root_compile (src, fname, batch)
68 from .Helpers
import ShutUp
as root_shutup
72 if 'dbg' in os.environ.get(
'CMTCONFIG',
'opt'):
73 compile_options +=
'g'
75 compile_options +=
'O'
81 src_file = tempfile.NamedTemporaryFile(prefix=
'root_aclic_',
83 src_file.write(textwrap.dedent(src).
encode())
97 _tempfiles.append (src_file)
101 import os.path
as osp
102 fname = osp.expanduser(osp.expandvars(fname))
105 assert os.access(fname, os.R_OK),
"could not read [%s]"%(fname,)
106 orig_root_lvl = ROOT.gErrorIgnoreLevel
107 ROOT.gErrorIgnoreLevel = ROOT.kWarning
110 sc = ROOT.gSystem.CompileMacro(fname, compile_options)
111 if sc == ROOT.kFALSE:
113 'problem compiling ROOT macro (rc=%s)'%(sc,)
116 ROOT.gErrorIgnoreLevel = orig_root_lvl
123 import PyUtils.Helpers
as H
124 with H.ShutUp(filters=[
126 'TClass::TClass:0: RuntimeWarning: no dictionary for.*'),
128 'Warning in <TEnvRec::ChangeValue>: duplicate entry.*'
131 cppyy.load_library(
"libRootUtilsPyROOTDict")
132 _ = root.RootUtils.PyBytes
134 read_root_file = root.RootUtils._pythonize_read_root_file
135 tell_root_file = root.RootUtils._pythonize_tell_root_file
137 def read(self, size=-1):
138 """read([size]) -> read at most size bytes, returned as a string.
140 If the size argument is negative or omitted, read until EOF is reached.
141 Notice that when in non-blocking mode, less data than what was requested
142 may be returned, even if no size parameter was given.
144 FIXME: probably doesn't follow python file-like conventions...
151 c_buf = read_root_file(self, size)
152 if c_buf
and c_buf.sz:
154 return bytes([ord(v[i])
for i
in range(v.size())])
161 c_buf = read_root_file(self, size)
162 if c_buf
and c_buf.sz:
164 chunk = bytes([ord(v[i])
for i
in range(v.size())])
170 root.TFile.read = read
173 root.TFile.seek = root.TFile.Seek
174 root.TFile.tell =
lambda self: tell_root_file(self)
185 tname = l.GetTypeName()
189 if tname
in [
'UInt_t',
'Int_t',
'ULong_t',
'Long_t',
'ULong64_t',
'Long64_t',
'UShort_t',
'Short_t',
'Bool_t']:
190 return tuple(l.GetValueLong64(i)
for i
in range(ndat))
191 elif tname
in [
'Float_t',
'Double_t',
'Float16_t',
'Double32_t']:
192 return tuple(l.GetValue(i)
for i
in range(ndat))
193 elif tname
in [
'UChar_t',
'Char_t']:
195 return l.GetValueString()
197 return tuple(l.GetValueLong64(i)
for i
in range(ndat))
199 if tname
in [
'UInt_t',
'Int_t',
'ULong_t',
'Long_t',
'ULong64_t',
'Long64_t',
'UShort_t',
'Short_t',
'Bool_t']:
200 return l.GetValueLong64()
201 elif tname
in [
'Float_t',
'Double_t',
'Float16_t',
'Double32_t']:
203 elif tname
in [
'UChar_t',
'Char_t']:
205 return l.GetValueString()
207 return l.GetValueLong64()
213 A helper class to dump in more or less human readable form the content of
218 object.__init__(self)
227 not isinstance(self.
root_file, ROOT.TFile)
or
229 raise IOError(
'could not open [%s]'% fname)
235 keys = [k.GetName()
for k
in self.
root_file.GetListOfKeys()]
238 if isinstance(o, ROOT.TTree):
247 from PyUtils.PoolFile
import PoolOpts
248 TTreeNames = PoolOpts.TTreeNames
249 RNTupleNames = PoolOpts.RNTupleNames
252 for name, klass
in ((TTreeNames.EventData, ROOT.TTree), (RNTupleNames.EventData, ROOT.RNTuple)):
253 if (obj := self.
root_file.
Get(name))
and isinstance(obj, klass):
257 raise AttributeError(
'No TTree named %r or RNTuple named %r in file %r' %
258 (TTreeNames.EventData, RNTupleNames.EventData,
262 not isinstance(obj, ROOT.TTree)
and not isinstance(obj, ROOT.RNTuple)):
263 raise AttributeError(
'No TTree or RNTuple named %r in file %r' %
267 if isinstance(obj, ROOT.RNTuple):
269 self.
obj = ROOT.RNTupleReader.Open(obj)
270 except AttributeError:
271 self.
obj = ROOT.Experimental.RNTupleReader.Open(obj)
272 elif isinstance(obj, ROOT.TTree):
277 def _dump(self, obj, itr_entries, leaves=None, retvecs=False, sortleaves=True):
280 RNTupleReader = ROOT.RNTupleReader
281 except AttributeError:
282 RNTupleReader = ROOT.Experimental.RNTupleReader
283 if isinstance(obj, ROOT.TTree):
284 yield from self.
_tree_dump(obj, itr_entries, leaves, retvecs, sortleaves)
285 elif isinstance(obj, RNTupleReader):
286 yield from self.
_reader_dump(obj, itr_entries, leaves, retvecs, sortleaves)
288 raise NotImplementedError(
"'_dump' not implemented for object of class=%r" %
289 (obj.__class__.__name__,))
291 def dump(self, tree_name, itr_entries, leaves=None, retvecs=False, sortleaves=True):
292 if (tree_name
is None and getattr(self,
"obj_name",
None)
is None or
293 tree_name
is not None and getattr(self,
"obj_name",
None) != tree_name):
295 yield from self.
_dump(self.
obj, itr_entries, leaves, retvecs, sortleaves)
297 def _tree_dump(self, tree, itr_entries, leaves=None, retvecs=False, sortleaves=True):
300 import AthenaPython.PyAthena
as PyAthena
301 _pythonize = PyAthena.RootUtils.PyROOTInspector.pyroot_inspect2
304 nentries = tree.GetEntries()
305 if leaves
is not None:
306 leaves = [
str(b).rstrip(
'\0')
for b
in leaves]
309 leaves =
sorted([b.GetName().rstrip(
'\0')
for b
in tree.GetListOfBranches()])
312 if isinstance(itr_entries, str):
313 if ':' in itr_entries:
321 from itertools
import islice
322 itr_entries = islice(
range(nentries),
323 *map(toint, itr_entries.split(
':')))
324 elif (
'range' in itr_entries
or
326 itr_entries = eval(itr_entries)
329 _n =
int(itr_entries)
330 itr_entries =
range(_n)
332 print (
"** err ** invalid 'itr_entries' argument. will iterate over all entries.")
333 itr_entries =
range(nentries)
334 elif isinstance(itr_entries, list):
335 itr_entries = itr_entries
337 itr_entries =
range(itr_entries)
342 isinstance_ = isinstance
344 for ientry
in itr_entries:
345 hdr =
":: entry [%05i]..." % (ientry,)
348 err = tree.LoadTree(ientry)
350 print (
"**err** loading tree for entry",ientry)
354 nbytes = tree.GetEntry(ientry)
356 print (
"**err** reading entry [%s] of tree [%s]" % (ientry, tree_name))
357 hdr =
":: entry [%05i]... [ERR]" % (ientry,)
362 for br_name
in leaves:
363 hdr =
":: branch [%s]..." % (br_name,)
368 br = tree.GetBranch (br_name)
369 if br.GetClassName() !=
'':
375 getattr (ROOT, br.GetClassName())
376 val = getattr(tree, br_name)
377 _vals += [ ([br_name], val) ]
379 for l
in br.GetListOfLeaves():
380 if (br.GetNleaves() == 1
and (br_name == l.GetName()
or
381 br_name.endswith(
'.' + l.GetName()))):
382 _vals += [ ([br_name], _getLeaf (l)) ]
384 _vals += [ ([br_name, l.GetName()], _getLeaf (l)) ]
385 for py_name, val
in _vals:
386 if val
is None:
continue
388 vals = _pythonize(val, py_name,
True, retvecs)
389 except Exception
as err:
390 print (
"**err** for branch [%s] val=%s (type=%s)" % (
391 br_name, val,
type(val),
396 viter =
sorted(vals, key =
lambda x:
'.'.
join(s
for s
in x[0]
if isinstance_(s, str_)))
400 n = list_(map_(str_, o[0]))
402 yield tree_name, ientry, n, v
408 def _reader_dump(self, reader, itr_entries, leaves=None, retvecs=False, sortleaves=True):
411 import AthenaPython.PyAthena
as PyAthena
412 _pythonize = PyAthena.RootUtils.PyROOTInspector.pyroot_inspect2
415 nentries = reader.GetNEntries()
416 from operator
import methodcaller
417 if leaves
is not None:
420 leaves =
sorted(map(methodcaller(
'GetFieldName'), reader.GetDescriptor().GetTopLevelFields()))
423 if isinstance(itr_entries, str):
424 if ':' in itr_entries:
432 from itertools
import islice
433 itr_entries = islice(
range(nentries),
434 *map(toint, itr_entries.split(
':')))
435 elif (
'range' in itr_entries
or
437 itr_entries = eval(itr_entries)
440 _n =
int(itr_entries)
441 itr_entries =
range(_n)
443 print (
"** err ** invalid 'itr_entries' argument. will iterate over all entries.")
444 itr_entries =
range(nentries)
445 elif isinstance(itr_entries, list):
446 itr_entries = itr_entries
448 itr_entries =
range(itr_entries)
453 isinstance_ = isinstance
456 from ROOT
import RException
458 from ROOT.Experimental
import RException
461 entry = reader.CreateEntry()
462 except AttributeError:
463 entry = reader.GetModel().CreateEntry()
465 for ientry
in itr_entries:
467 reader.LoadEntry(ientry, entry)
468 except RException
as err:
469 from traceback
import format_exception
471 print(
"Exception reading entry=%05i of ntuple %r\n%s" %
472 (ientry, ntuple_name,
"".
join(format_exception(err))), file=sys.stderr)
476 for br_name
in leaves:
478 token = entry.GetToken(br_name)
479 typeName = entry.GetTypeName(token)
485 getattr(ROOT, typeName)
489 vals = _pythonize(val, py_name,
True, retvecs)
490 except Exception
as err:
491 print(
"**err** for branch [%s] val=%s (type=%s)" %
492 (br_name, val,
type(val)))
496 viter =
sorted(vals, key =
lambda x:
'.'.
join(s
for s
in x[0]
if isinstance_(s, str_)))
500 n = list_(map_(str_, o[0]))
502 yield ntuple_name, ientry, n, v
508 def no_raise(msg, fct, *args, **kwds):
513 except Exception
as xerr:
516 assert not caught,
"%s:\n%s\nERROR" % (msg, err,)
518 no_raise(
"problem pythonizing TFile", fct=_pythonize_tfile)
519 no_raise(
"problem compiling dummy one-liner",
520 root_compile,
"void foo1() { return ; }")
521 no_raise(
"problem compiling dummy one-liner w/ kwds",
522 fct=root_compile, src=
"void foo1a() { return ; }")
525 dummy = tempfile.NamedTemporaryFile(prefix=
"foo_",suffix=
".cxx")
526 with tempfile.NamedTemporaryFile(prefix=
"foo_",suffix=
".cxx")
as tmp:
527 tmp.write (b
"void foo2() { return ; }\n")
529 no_raise(
"problem compiling a file",
530 fct=root_compile, fname=tmp.name)
535 if __name__ ==
"__main__":