Belle II Software  release-05-02-19
root_utils.py
1 """Various functions to interact with ROOT objects and the runtime environment"""
2 
3 import ROOT
4 import contextlib
5 from functools import singledispatch
6 
7 
8 def root_walk(tdirectory):
9  """Walks the content of a TDirectory similar to os.walk.
10 
11  Yields 3-tuples of current TDirectories, contained TObjects and contained TDirectories
12  for each of the directories nested inside the given TDirectory in a depth first manner.
13 
14  Yields
15  ------
16  (TDirectory, list(TObject), list(TDirectory))
17  """
18  tkeys = tdirectory.GetListOfKeys()
19 
20  tobjects = []
21  tdirectories = []
22  for tkey in tkeys:
23  tobject = tdirectory.Get(tkey.GetName())
24  if isinstance(tobject, ROOT.TDirectory):
25  tdirectories.append(tobject)
26  else:
27  tobjects.append(tobject)
28 
29  yield tdirectory, tobjects, tdirectories
30 
31  # Continue depth first
32  for sub_tdirectory in tdirectories:
33  for tdirectory, tobjects, tdirectories in root_walk(sub_tdirectory):
34  yield tdirectory, tobjects, tdirectories
35 
36 
37 @contextlib.contextmanager
38 def root_open(tfile_or_file_path, tfile_options=None):
39  """Context manager to open a TFile.
40 
41  If a file path is given open the TFile and close it after the context is left.
42  If an already opened TFile is received simply return it and do not close on exit.
43 
44  Parameters
45  ----------
46  tfile_or_file_path : str or ROOT.TFile
47  Path to the file or the TFile that should be activated.
48  tfile_options : str
49  Option string forwarded to the ROOT.TFile constructor
50  Typical options as "RECREATE", "READ" or "UPDATE".
51  """
52  if isinstance(tfile_or_file_path, ROOT.TFile):
53  tfile = tfile_or_file_path
54  with root_cd(tfile) as tfile:
55  yield tfile
56  else:
57  save_tdirectory = ROOT.gROOT.CurrentDirectory()
58 
59  if tfile_options is None:
60  tfile = ROOT.TFile(tfile_or_file_path)
61  else:
62  tfile = ROOT.TFile(tfile_or_file_path, tfile_options)
63 
64  try:
65  yield tfile
66  finally:
67  tfile.Close()
68  save_tdirectory.cd()
69 
70 
71 @contextlib.contextmanager
72 def root_cd(tdirectory):
73  """Context manager that temporarily switches the current global ROOT directory while in the context.
74 
75  If a string as the name of a directory is given as the argument
76  try to switch to the directory with that name in the current ROOT folder.
77 
78  If it is not present create it.
79 
80  Parameters
81  ----------
82  tdirectory : ROOT.TDirectory or str
83  ROOT directory to switch to or name of a folder to switch.
84 
85  Returns
86  -------
87  ROOT.TDirectory
88  The new current ROOT directory.
89  """
90 
91  # Do not use ROOT.gDirectory here.
92  # Since ROOT.gDirectory gets transported as a reference it changes on a call to cd() as well,
93  # and can therefore not serve to save the former directory.
94  save_tdirectory = ROOT.gROOT.CurrentDirectory()
95 
96  if not tdirectory or "." == tdirectory:
97  tdirectory = save_tdirectory
98 
99  elif isinstance(tdirectory, str):
100  tdirectory_name = tdirectory
101 
102  # Look for the tdirectory with the name
103  # before trying to create it
104  tdirectory = save_tdirectory.GetDirectory(tdirectory_name)
105  if not tdirectory:
106  tdirectory = save_tdirectory.mkdir(tdirectory_name, tdirectory_name)
107  if not tdirectory:
108  raise RuntimeError("Could not create or find folder %s" % tdirectory_name)
109 
110  # If tdirectory_name is as hierachy like a/b/c make sure we select the most nested folder
111  # (and not a as the documentation of TDirectory.mkdir suggests).
112  tdirectory = save_tdirectory.GetDirectory(tdirectory_name)
113 
114  try:
115  if tdirectory is not None:
116  tdirectory.cd()
117  yield tdirectory
118 
119  finally:
120  save_tdirectory.cd()
121 
122 
123 def root_save_name(name):
124  """Strips all meta characters that might be unsafe to use as a ROOT name.
125 
126  Parameters
127  ----------
128  name : str
129  A name that should be transformed
130 
131  Returns
132  -------
133  str
134  Name with potentially harmful characters deleted / replaced.
135  """
136  deletechars = str.maketrans("", "", r"/$\#{}()[]=")
137  name = name.replace(' ', '_').replace('-', '_').replace(',', '_').translate(deletechars)
138  return name
139 
140 
141 def root_browse(tobject):
142  """Open a browser and show the given object.
143 
144  Parameters
145  ----------
146  tobject : ROOT.TObject
147  The object to be shown
148 
149  Returns
150  -------
151  ROOT.TBrowser
152  The new TBrowser used to show the object.
153  """
154  # Set the style to the style desired by the validation plots.
157 
158  tbrowser = ROOT.TBrowser()
159  if isinstance(tobject, ROOT.TObject):
160  if hasattr(tobject, "Browse"):
161  tobject.Browse(tbrowser)
162  else:
163  tbrowser.BrowseObject(tobject)
164  else:
165  raise ValueError("Can only browse ROOT objects inheriting from TObject.")
166  return tbrowser
167 
168 
169 @singledispatch
170 def root_ls(obj):
171  """Returns a list of names that are contained in the given obj.
172 
173  This is a convinience function to invesitigate the content of ROOT objects,
174  that dispatches on to object type and retieves different things depending on the type.
175  If the obj is a string it is interpreted as a filename.
176  """
177  return list(obj)
178 
179 
180 @root_ls.register(str)
181 def __(filename):
182  """Overloaded function for root_ls for filenames (e.g. opens the file and ls its content)."""
183  rootFile = ROOT.TFile(filename)
184  result = ls(rootInputFile)
185 
186  rootFile.Close()
187  del rootFile
188  return result
189 
190 
191 @root_ls.register(ROOT.TDirectory)
192 def __(tDirectory):
193  """Overloaded function for root_ls for ROOT directories (e.g. list the keys in the directory)."""
194  tKeys = list(tDirectory.GetListOfKeys())
195  result = sorted([tKey.GetName() for tKey in tKeys])
196  return result
197 
198 
199 @root_ls.register(ROOT.TTree)
200 @root_ls.register(ROOT.TNtuple)
201 def __(tTree):
202  """Overloaded function for root_ls for trees and ntuples (e.g. list the keys in the tuple/tree)."""
203  tBranches = list(tTree.GetListOfBranches())
204  result = sorted([tBranch.GetName() for tBranch in tBranches])
205  return result
tracking.root_utils.root_ls
def root_ls(obj)
Definition: root_utils.py:170
tracking.root_utils.root_save_name
def root_save_name(name)
Definition: root_utils.py:123
tracking.root_utils.root_walk
def root_walk(tdirectory)
Definition: root_utils.py:8
tracking.validation.plot.ValidationPlot.set_tstyle
def set_tstyle(cls)
Definition: plot.py:2263
tracking.root_utils.__
def __(filename)
Definition: root_utils.py:181
tracking.root_utils.root_open
def root_open(tfile_or_file_path, tfile_options=None)
Definition: root_utils.py:38
tracking.root_utils.root_browse
def root_browse(tobject)
Definition: root_utils.py:141
tracking.validation.plot
Definition: plot.py:1
tracking.root_utils.root_cd
def root_cd(tdirectory)
Definition: root_utils.py:72