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 tname in ['UInt_t', 'Int_t', 'ULong64_t', 'Long64_t']:
188  return [l.GetValueLong64(i) for i in range(ndat)]
189  if tname in ['Float_t', 'Double_t']:
190  return [l.GetValue(i) for i in range(ndat)]
191  if tname in ['Char_t']:
192  try:
193  return l.GetValueString() # TLeafC for variable size string
194  except Exception:
195  return [l.GetValue(i) for i in range(ndat)] # TLeafB for 8-bit integers
196  return None
197 
199  """
200  A helper class to dump in more or less human readable form the content of
201  any TTree.
202  """
203 
204  def __init__(self, fname, tree_name="CollectionTree"):
205  object.__init__(self)
206 
207  ROOT = import_root()
208 
209  # remember if an error or problem occurred during the dump
210  self.allgood = True
211 
212  self.root_file = ROOT.TFile.Open(fname)
213  if (self.root_file is None or
214  not isinstance(self.root_file, ROOT.TFile) or
215  not self.root_file.IsOpen()):
216  raise IOError('could not open [%s]'% fname)
217 
218  self.tree = self.root_file.Get(tree_name)
219  if self.tree is None or not isinstance(self.tree, ROOT.TTree):
220  raise AttributeError('no tree [%s] in file [%s]', tree_name, fname)
221 
222  if 0:
223  self._trees = []
224  keys = [k.GetName() for k in self.root_file.GetListOfKeys()]
225  for k in keys:
226  o = self.root_file.Get(k)
227  if isinstance(o, ROOT.TTree):
228  self._trees.append(k)
229  pass
230 
231  return
232 
233  def dump(self, tree_name, itr_entries, leaves=None, retvecs=False, sortleaves=True):
234 
235  ROOT = import_root()
236  import AthenaPython.PyAthena as PyAthena
237  _pythonize = PyAthena.RootUtils.PyROOTInspector.pyroot_inspect2
238 
239  self.tree = self.root_file.Get(tree_name)
240  if self.tree is None or not isinstance(self.tree, ROOT.TTree):
241  raise AttributeError('no tree [%s] in file [%s]', tree_name, self.root_file.GetName())
242 
243  tree = self.tree
244  nentries = tree.GetEntries()
245  branches = sorted([b.GetName().rstrip('\0') for b in tree.GetListOfBranches()])
246  if leaves is None: leaves = branches
247  else: leaves = [str(b).rstrip('\0') for b in leaves]
248 
249  # handle itr_entries
250  if isinstance(itr_entries, str):
251  if ':' in itr_entries:
252  def toint(s):
253  if s == '':
254  return None
255  try:
256  return int(s)
257  except ValueError:
258  return s
259  from itertools import islice
260  itr_entries = islice(range(nentries),
261  *map(toint, itr_entries.split(':')))
262  elif ('range' in itr_entries or
263  ',' in itr_entries):
264  itr_entries = eval(itr_entries)
265  else:
266  try:
267  _n = int(itr_entries)
268  itr_entries = range(_n)
269  except ValueError:
270  print ("** err ** invalid 'itr_entries' argument. will iterate over all entries.")
271  itr_entries = range(nentries)
272  elif isinstance(itr_entries, list):
273  itr_entries = itr_entries
274  else:
275  itr_entries = range(itr_entries)
276 
277  list_ = list
278  map_ = map
279  str_ = str
280  isinstance_ = isinstance
281 
282  for ientry in itr_entries:
283  hdr = ":: entry [%05i]..." % (ientry,)
284  #print (hdr)
285  #print (hdr, file=self.fout)
286  err = tree.LoadTree(ientry)
287  if err < 0:
288  print ("**err** loading tree for entry",ientry)
289  self.allgood = False
290  break
291 
292  nbytes = tree.GetEntry(ientry)
293  if nbytes <= 0:
294  print ("**err** reading entry [%s] of tree [%s]" % (ientry, tree_name))
295  hdr = ":: entry [%05i]... [ERR]" % (ientry,)
296  print (hdr)
297  self.allgood = False
298  continue
299 
300  for br_name in leaves:
301  hdr = ":: branch [%s]..." % (br_name,)
302  #print (hdr)
303  #tree.GetBranch(br_name).GetEntry(ientry)
304  py_name = [br_name]
305 
306  br = tree.GetBranch (br_name)
307  if br.GetClassName() != '':
308  # Make sure dictionaries are completely loaded before
309  # trying to fetch it from ROOT. Otherwise we can run
310  # into cling parse failures due to it synthesizing
311  # incorrect forward declarations.
312  # See ATEAM-1000.
313  getattr (ROOT, br.GetClassName())
314  val = getattr(tree, br_name)
315  else:
316  vals = [_getLeaf (l) for l in br.GetListOfLeaves()]
317  if len(vals) == 0:
318  val = None
319  elif len(vals) == 1:
320  val = vals
321  else:
322  val = tuple(vals)
323  if not (val is None):
324  #print ("-->",val,br_name)
325  try:
326  vals = _pythonize(val, py_name, True, retvecs)
327  except Exception as err:
328  print ("**err** for branch [%s] val=%s (type=%s)" % (
329  br_name, val, type(val),
330  ))
331  self.allgood = False
332  print (err)
333  if sortleaves:
334  viter = sorted(vals, key = lambda x: '.'.join(s for s in x[0] if isinstance_(s, str_)))
335  else:
336  viter = vals
337  for o in viter:
338  n = list_(map_(str_, o[0]))
339  v = o[1]
340  yield tree_name, ientry, n, v
341 
342  pass # loop over branch names
343  pass # loop over entries
344  pass # class RootFileDumper
345 
346 
348  root = import_root() # noqa: F841
349  def no_raise(msg, fct, *args, **kwds):
350  caught = False
351  err = None
352  try:
353  fct(*args, **kwds)
354  except Exception as xerr:
355  err = xerr
356  caught = True
357  assert not caught, "%s:\n%s\nERROR" % (msg, err,)
358 
359  no_raise("problem pythonizing TFile", fct=_pythonize_tfile)
360  no_raise("problem compiling dummy one-liner",
361  root_compile, "void foo1() { return ; }")
362  no_raise("problem compiling dummy one-liner w/ kwds",
363  fct=root_compile, src="void foo1a() { return ; }")
364  import tempfile
365  # PvG workaround for ROOT-7059
366  dummy = tempfile.NamedTemporaryFile(prefix="foo_",suffix=".cxx") # noqa: F841
367  with tempfile.NamedTemporaryFile(prefix="foo_",suffix=".cxx") as tmp:
368  tmp.write (b"void foo2() { return ; }\n")
369  tmp.flush()
370  no_raise("problem compiling a file",
371  fct=root_compile, fname=tmp.name)
372 
373  print ("OK")
374  return True
375 
376 if __name__ == "__main__":
377  _test_main()
378 
read
IovVectorMap_t read(const Folder &theFolder, const SelectionCriterion &choice, const unsigned int limit=10)
Definition: openCoraCool.cxx:569
python.RootUtils.RootFileDumper._trees
_trees
Definition: RootUtils.py:223
python.RootUtils.RootFileDumper
Definition: RootUtils.py:198
CaloCellPos2Ntuple.int
int
Definition: CaloCellPos2Ntuple.py:24
python.RootUtils._getLeaf
def _getLeaf(l)
Definition: RootUtils.py:184
python.AthDsoLogger.fct
fct
Definition: AthDsoLogger.py:43
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:233
python.RootUtils._root_compile
def _root_compile(src, fname, batch)
Definition: RootUtils.py:66
python.RootUtils.import_root
def import_root(batch=True)
functions --------------------------------------------------------------—
Definition: RootUtils.py:22
python.RootUtils.RootFileDumper.root_file
root_file
Definition: RootUtils.py:212
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:179
AtlasMcWeight::encode
number_type encode(double weight)
Definition: AtlasMcWeight.cxx:65
plotBeamSpotVxVal.range
range
Definition: plotBeamSpotVxVal.py:195
python.RootUtils._test_main
def _test_main()
Definition: RootUtils.py:347
DerivationFramework::TriggerMatchingUtils::sorted
std::vector< typename T::value_type > sorted(T begin, T end)
Helper function to create a sorted vector from an unsorted one.
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.tree
tree
Definition: RootUtils.py:218
python.RootUtils.RootFileDumper.__init__
def __init__(self, fname, tree_name="CollectionTree")
Definition: RootUtils.py:204
python.CaloScaleNoiseConfig.type
type
Definition: CaloScaleNoiseConfig.py:78
pickleTool.object
object
Definition: pickleTool.py:30
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:210