ATLAS Offline Software
RootUtils.py
Go to the documentation of this file.
1 # Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
2 
3 # @file PyUtils.RootUtils
4 # @author Sebastien Binet
5 # @purpose a few utils to ease the day-to-day work with ROOT
6 # @date November 2009
7 
8 __doc__ = "a few utils to ease the day-to-day work with ROOT"
9 __author__ = "Sebastien Binet"
10 
11 __all__ = [
12  'import_root',
13  'root_compile',
14  ]
15 
16 
17 import os
18 import re
19 from functools import cache
20 
21 
22 def import_root(batch=True):
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.)
26 
27  e.g.
28  >>> ROOT = import_root(batch=True)
29  >>> f = ROOT.TFile.Open(...)
30  """
31  import ROOT
32  ROOT.gROOT.SetBatch(batch)
33  if batch:
34  ROOT.PyConfig.IgnoreCommandLineOptions = True
35  import cppyy # noqa: F401
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)
41  return ROOT
42 
43 _tempfiles = []
44 _first_compile = True
45 def root_compile(src=None, fname=None, batch=True):
46  """a helper method to compile a set of C++ statements (via ``src``) or
47  a C++ file (via ``fname``) via ACLiC
48  """
49  if src is not None and fname is not None:
50  raise ValueError("'src' xor 'fname' should be not None, *not* both")
51 
52  if src is None and fname is None:
53  raise ValueError("'src' xor 'fname' should be None, *not* both")
54 
55  # Cling bug workaround: Cling will try to find a definition for the
56  # hidden __gmon_start__ by opening all libraries on LD_LIBRARY_PATH.
57  # But it will crash if it encounters a separate-debug library.
58  # Work around by adding a dummy definition of __gmon_start__.
59  # See !31633.
60  global _first_compile
61  if _first_compile:
62  _first_compile = False
63  root_compile ('extern "C" { void __gmon_start__(){}; }', None, True)
64  return _root_compile (src, fname, batch)
65 
66 def _root_compile (src, fname, batch):
67  import os
68  from .Helpers import ShutUp as root_shutup
69 
70  ROOT = import_root(batch=batch)
71  compile_options = "f"
72  if 'dbg' in os.environ.get('CMTCONFIG', 'opt'):
73  compile_options += 'g'
74  else:
75  compile_options += 'O'
76 
77  src_file = None
78  if src:
79  import textwrap
80  import tempfile
81  src_file = tempfile.NamedTemporaryFile(prefix='root_aclic_',
82  suffix='.cxx')
83  src_file.write(textwrap.dedent(src).encode())
84  src_file.flush()
85  src_file.seek(0)
86  fname = src_file.name
87 
88  # Apparently, cling caches files by inode.
89  # If you ask it to read a file that has the same inode as one
90  # that it has already read, then it will just use the cached
91  # contents rather than rereading. This, however, doesn't play
92  # very well if we're reading temp files, where inodes may be reused,
93  # giving rise to hard-to-reproduce failures.
94  #
95  # Try to avoid this by keeping the temp files open until the
96  # the program exits.
97  _tempfiles.append (src_file)
98  pass
99 
100  elif fname:
101  import os.path as osp
102  fname = osp.expanduser(osp.expandvars(fname))
103  pass
104 
105  assert os.access(fname, os.R_OK), "could not read [%s]"%(fname,)
106  orig_root_lvl = ROOT.gErrorIgnoreLevel
107  ROOT.gErrorIgnoreLevel = ROOT.kWarning
108  try:
109  with root_shutup():
110  sc = ROOT.gSystem.CompileMacro(fname, compile_options)
111  if sc == ROOT.kFALSE:
112  raise RuntimeError(
113  'problem compiling ROOT macro (rc=%s)'%(sc,)
114  )
115  finally:
116  ROOT.gErrorIgnoreLevel = orig_root_lvl
117  return
118 
119 @cache
121  import cppyy
122  root = import_root()
123  import PyUtils.Helpers as H
124  with H.ShutUp(filters=[
125  re.compile(
126  'TClass::TClass:0: RuntimeWarning: no dictionary for.*'),
127  re.compile(
128  'Warning in <TEnvRec::ChangeValue>: duplicate entry.*'
129  ),
130  ]):
131  cppyy.load_library("libRootUtilsPyROOTDict")
132  _ = root.RootUtils.PyBytes
133  #MN: lines below fail in ROOT6 if PCM from RootUtils is not found
134  read_root_file = root.RootUtils._pythonize_read_root_file
135  tell_root_file = root.RootUtils._pythonize_tell_root_file
136  pass
137  def read(self, size=-1):
138  """read([size]) -> read at most size bytes, returned as a string.
139 
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.
143 
144  FIXME: probably doesn't follow python file-like conventions...
145  """
146  SZ = 4096
147 
148  if size>=0:
149  #size = _adjust_sz(size)
150  #print ("-->0",self.tell(),size)
151  c_buf = read_root_file(self, size)
152  if c_buf and c_buf.sz:
153  v = c_buf.buf
154  return bytes([ord(v[i]) for i in range(v.size())])
155  return ''
156  else:
157  size = SZ
158  out = []
159  while True:
160  #size = _adjust_sz(size)
161  c_buf = read_root_file(self, size)
162  if c_buf and c_buf.sz:
163  v = c_buf.buf
164  chunk = bytes([ord(v[i]) for i in range(v.size())])
165  out.append(chunk)
166  else:
167  break
168  return b''.join(out)
169 
170  root.TFile.read = read
171  del read
172 
173  root.TFile.seek = root.TFile.Seek
174  root.TFile.tell = lambda self: tell_root_file(self)
175 
181  return
182 
183 
184 def _getLeaf (l):
185  tname = l.GetTypeName()
186  ndat = l.GetNdata()
187  if (l.GetLeafCount() # a varying length array
188  or ndat > 1): # a fixed size array
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']:
194  try:
195  return l.GetValueString() # TLeafC for variable size string
196  except Exception:
197  return tuple(l.GetValueLong64(i) for i in range(ndat)) # TLeafB for 8-bit integers
198  elif ndat == 1: # a single value
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']:
202  return l.GetValue()
203  elif tname in ['UChar_t', 'Char_t']:
204  try:
205  return l.GetValueString() # TLeafC for variable size string
206  except Exception:
207  return l.GetValueLong64() # TLeafB for 8-bit integers
208 
209  return None
210 
212  """
213  A helper class to dump in more or less human readable form the content of
214  any TTree.
215  """
216 
217  def __init__(self, fname, tree_name=None):
218  object.__init__(self)
219 
220  ROOT = import_root()
221 
222  # remember if an error or problem occurred during the dump
223  self.allgood = True
224 
225  self.root_file = ROOT.TFile.Open(fname)
226  if (self.root_file is None or
227  not isinstance(self.root_file, ROOT.TFile) or
228  not self.root_file.IsOpen()):
229  raise IOError('could not open [%s]'% fname)
230 
231  self.__init_obj(tree_name)
232 
233  if 0:
234  self._trees = []
235  keys = [k.GetName() for k in self.root_file.GetListOfKeys()]
236  for k in keys:
237  o = self.root_file.Get(k)
238  if isinstance(o, ROOT.TTree):
239  self._trees.append(k)
240  pass
241 
242  return
243 
244  def __init_obj(self, obj_name):
245 
246  ROOT = import_root()
247  from PyUtils.PoolFile import PoolOpts
248  TTreeNames = PoolOpts.TTreeNames
249  RNTupleNames = PoolOpts.RNTupleNames
250 
251  if obj_name is None:
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):
254  self.obj_name = name
255  break
256  else:
257  raise AttributeError('No TTree named %r or RNTuple named %r in file %r' %
258  (TTreeNames.EventData, RNTupleNames.EventData,
259  self.root_file.GetName()))
260  else:
261  if (not (obj := self.root_file.Get(obj_name)) or
262  not isinstance(obj, ROOT.TTree) and not isinstance(obj, ROOT.RNTuple)):
263  raise AttributeError('No TTree or RNTuple named %r in file %r' %
264  (obj_name, self.root_file.GetName()))
265  self.obj_name = obj_name
266 
267  if isinstance(obj, ROOT.RNTuple):
268  try:
269  self.obj = ROOT.RNTupleReader.Open(obj)
270  except AttributeError:
271  self.obj = ROOT.Experimental.RNTupleReader.Open(obj)
272  elif isinstance(obj, ROOT.TTree):
273  self.obj = obj
274  # in case it is used somewhere
275  self.tree = self.obj
276 
277  def _dump(self, obj, itr_entries, leaves=None, retvecs=False, sortleaves=True):
278  ROOT = import_root()
279  try:
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)
287  else:
288  raise NotImplementedError("'_dump' not implemented for object of class=%r" %
289  (obj.__class__.__name__,))
290 
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):
294  self.__init_obj(tree_name)
295  yield from self._dump(self.obj, itr_entries, leaves, retvecs, sortleaves)
296 
297  def _tree_dump(self, tree, itr_entries, leaves=None, retvecs=False, sortleaves=True):
298 
299  ROOT = import_root()
300  import AthenaPython.PyAthena as PyAthena
301  _pythonize = PyAthena.RootUtils.PyROOTInspector.pyroot_inspect2
302 
303  tree_name = self.obj_name
304  nentries = tree.GetEntries()
305  if leaves is not None:
306  leaves = [str(b).rstrip('\0') for b in leaves]
307  leaves.sort()
308  else:
309  leaves = sorted([b.GetName().rstrip('\0') for b in tree.GetListOfBranches()])
310 
311  # handle itr_entries
312  if isinstance(itr_entries, str):
313  if ':' in itr_entries:
314  def toint(s):
315  if s == '':
316  return None
317  try:
318  return int(s)
319  except ValueError:
320  return s
321  from itertools import islice
322  itr_entries = islice(range(nentries),
323  *map(toint, itr_entries.split(':')))
324  elif ('range' in itr_entries or
325  ',' in itr_entries):
326  itr_entries = eval(itr_entries)
327  else:
328  try:
329  _n = int(itr_entries)
330  itr_entries = range(_n)
331  except ValueError:
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
336  else:
337  itr_entries = range(itr_entries)
338 
339  list_ = list
340  map_ = map
341  str_ = str
342  isinstance_ = isinstance
343 
344  for ientry in itr_entries:
345  hdr = ":: entry [%05i]..." % (ientry,)
346  #print (hdr)
347  #print (hdr, file=self.fout)
348  err = tree.LoadTree(ientry)
349  if err < 0:
350  print ("**err** loading tree for entry",ientry)
351  self.allgood = False
352  break
353 
354  nbytes = tree.GetEntry(ientry)
355  if nbytes <= 0:
356  print ("**err** reading entry [%s] of tree [%s]" % (ientry, tree_name))
357  hdr = ":: entry [%05i]... [ERR]" % (ientry,)
358  print (hdr)
359  self.allgood = False
360  continue
361 
362  for br_name in leaves:
363  hdr = ":: branch [%s]..." % (br_name,)
364  #print (hdr)
365  #tree.GetBranch(br_name).GetEntry(ientry)
366  _vals = list()
367 
368  br = tree.GetBranch (br_name)
369  if br.GetClassName() != '':
370  # Make sure dictionaries are completely loaded before
371  # trying to fetch it from ROOT. Otherwise we can run
372  # into cling parse failures due to it synthesizing
373  # incorrect forward declarations.
374  # See ATEAM-1000.
375  getattr (ROOT, br.GetClassName())
376  val = getattr(tree, br_name)
377  _vals += [ ([br_name], val) ]
378  else:
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)) ]
383  else:
384  _vals += [ ([br_name, l.GetName()], _getLeaf (l)) ]
385  for py_name, val in _vals:
386  if val is None: continue
387  try:
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),
392  ))
393  self.allgood = False
394  print (err)
395  if sortleaves:
396  viter = sorted(vals, key = lambda x: '.'.join(s for s in x[0] if isinstance_(s, str_)))
397  else:
398  viter = vals
399  for o in viter:
400  n = list_(map_(str_, o[0]))
401  v = o[1]
402  yield tree_name, ientry, n, v
403 
404  pass # loop over branch names
405  pass # loop over entries
406  pass # class RootFileDumper
407 
408  def _reader_dump(self, reader, itr_entries, leaves=None, retvecs=False, sortleaves=True):
409 
410  ROOT = import_root()
411  import AthenaPython.PyAthena as PyAthena
412  _pythonize = PyAthena.RootUtils.PyROOTInspector.pyroot_inspect2
413 
414  ntuple_name = self.obj_name
415  nentries = reader.GetNEntries()
416  from operator import methodcaller
417  if leaves is not None:
418  leaves = sorted(leaves)
419  else:
420  leaves = sorted(map(methodcaller('GetFieldName'), reader.GetDescriptor().GetTopLevelFields()))
421 
422  # handle itr_entries
423  if isinstance(itr_entries, str):
424  if ':' in itr_entries:
425  def toint(s):
426  if s == '':
427  return None
428  try:
429  return int(s)
430  except ValueError:
431  return s
432  from itertools import islice
433  itr_entries = islice(range(nentries),
434  *map(toint, itr_entries.split(':')))
435  elif ('range' in itr_entries or
436  ',' in itr_entries):
437  itr_entries = eval(itr_entries)
438  else:
439  try:
440  _n = int(itr_entries)
441  itr_entries = range(_n)
442  except ValueError:
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
447  else:
448  itr_entries = range(itr_entries)
449 
450  list_ = list
451  map_ = map
452  str_ = str
453  isinstance_ = isinstance
454 
455  try:
456  from ROOT import RException
457  except ImportError:
458  from ROOT.Experimental import RException
459 
460  try:
461  entry = reader.CreateEntry()
462  except AttributeError:
463  entry = reader.GetModel().CreateEntry()
464 
465  for ientry in itr_entries:
466  try:
467  reader.LoadEntry(ientry, entry)
468  except RException as err:
469  from traceback import format_exception
470  import sys
471  print("Exception reading entry=%05i of ntuple %r\n%s" %
472  (ientry, ntuple_name, "".join(format_exception(err))), file=sys.stderr)
473  self.allgood = False
474  continue
475 
476  for br_name in leaves:
477  py_name = [br_name]
478  token = entry.GetToken(br_name)
479  typeName = entry.GetTypeName(token)
480  # Make sure dictionaries are completely loaded before
481  # trying to fetch it from ROOT. Otherwise we can run
482  # into cling parse failures due to it synthesizing
483  # incorrect forward declarations.
484  # See ATEAM-1000.
485  getattr(ROOT, typeName)
486  val = entry[token]
487  if val is not None:
488  try:
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)))
493  self.allgood = False
494  print(err)
495  if sortleaves:
496  viter = sorted(vals, key = lambda x: '.'.join(s for s in x[0] if isinstance_(s, str_)))
497  else:
498  viter = vals
499  for o in viter:
500  n = list_(map_(str_, o[0]))
501  v = o[1]
502  yield ntuple_name, ientry, n, v
503 
504 
505 
507  root = import_root() # noqa: F841
508  def no_raise(msg, fct, *args, **kwds):
509  caught = False
510  err = None
511  try:
512  fct(*args, **kwds)
513  except Exception as xerr:
514  err = xerr
515  caught = True
516  assert not caught, "%s:\n%s\nERROR" % (msg, err,)
517 
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 ; }")
523  import tempfile
524  # PvG workaround for ROOT-7059
525  dummy = tempfile.NamedTemporaryFile(prefix="foo_",suffix=".cxx") # noqa: F841
526  with tempfile.NamedTemporaryFile(prefix="foo_",suffix=".cxx") as tmp:
527  tmp.write (b"void foo2() { return ; }\n")
528  tmp.flush()
529  no_raise("problem compiling a file",
530  fct=root_compile, fname=tmp.name)
531 
532  print ("OK")
533  return True
534 
535 if __name__ == "__main__":
536  _test_main()
537 
read
IovVectorMap_t read(const Folder &theFolder, const SelectionCriterion &choice, const unsigned int limit=10)
Definition: openCoraCool.cxx:569
DerivationFramework::TriggerMatchingUtils::sorted
std::vector< typename R::value_type > sorted(const R &r, PROJ proj={})
Helper function to create a sorted vector from an unsorted range.
python.RootUtils.RootFileDumper._trees
_trees
Definition: RootUtils.py:234
python.RootUtils.RootFileDumper
Definition: RootUtils.py:211
python.RootUtils._getLeaf
def _getLeaf(l)
Definition: RootUtils.py:184
python.RootUtils.RootFileDumper._tree_dump
def _tree_dump(self, tree, itr_entries, leaves=None, retvecs=False, sortleaves=True)
Definition: RootUtils.py:297
python.AthDsoLogger.fct
fct
Definition: AthDsoLogger.py:42
dumpHVPathFromNtuple.append
bool append
Definition: dumpHVPathFromNtuple.py:91
python.RootUtils.RootFileDumper.dump
def dump(self, tree_name, itr_entries, leaves=None, retvecs=False, sortleaves=True)
Definition: RootUtils.py:291
python.RootUtils._root_compile
def _root_compile(src, fname, batch)
Definition: RootUtils.py:66
python.RootUtils.RootFileDumper._dump
def _dump(self, obj, itr_entries, leaves=None, retvecs=False, sortleaves=True)
Definition: RootUtils.py:277
python.CaloAddPedShiftConfig.type
type
Definition: CaloAddPedShiftConfig.py:42
python.RootUtils.import_root
def import_root(batch=True)
functions --------------------------------------------------------------—
Definition: RootUtils.py:22
python.RootUtils.RootFileDumper.root_file
root_file
Definition: RootUtils.py:225
Get
T * Get(TFile &f, const std::string &n, const std::string &dir="", const chainmap_t *chainmap=0, std::vector< std::string > *saved=0)
get a histogram given a path, and an optional initial directory if histogram is not found,...
Definition: comparitor.cxx:181
AtlasMcWeight::encode
number_type encode(double weight)
Definition: AtlasMcWeight.cxx:65
python.RootUtils.RootFileDumper._reader_dump
def _reader_dump(self, reader, itr_entries, leaves=None, retvecs=False, sortleaves=True)
Definition: RootUtils.py:408
python.RootUtils.RootFileDumper.__init_obj
def __init_obj(self, obj_name)
Definition: RootUtils.py:244
plotBeamSpotVxVal.range
range
Definition: plotBeamSpotVxVal.py:194
python.RootUtils._test_main
def _test_main()
Definition: RootUtils.py:506
histSizes.list
def list(name, path='/')
Definition: histSizes.py:38
print
void print(char *figname, TCanvas *c1)
Definition: TRTCalib_StrawStatusPlots.cxx:25
TCS::join
std::string join(const std::vector< std::string > &v, const char c=',')
Definition: Trigger/TrigT1/L1Topo/L1TopoCommon/Root/StringUtils.cxx:10
python.RootUtils.RootFileDumper.obj_name
obj_name
Definition: RootUtils.py:254
python.RootUtils.RootFileDumper.obj
obj
Definition: RootUtils.py:269
python.CaloAddPedShiftConfig.int
int
Definition: CaloAddPedShiftConfig.py:45
python.RootUtils.RootFileDumper.tree
tree
Definition: RootUtils.py:275
pickleTool.object
object
Definition: pickleTool.py:29
python.RootUtils.RootFileDumper.__init__
def __init__(self, fname, tree_name=None)
Definition: RootUtils.py:217
str
Definition: BTagTrackIpAccessor.cxx:11
python.RootUtils.root_compile
def root_compile(src=None, fname=None, batch=True)
Definition: RootUtils.py:45
python.RootUtils._pythonize_tfile
def _pythonize_tfile()
Definition: RootUtils.py:120
python.RootUtils.RootFileDumper.allgood
allgood
Definition: RootUtils.py:223