1 """Various functions to interact with ROOT objects and the runtime environment"""
5 from functools
import singledispatch
9 """Walks the content of a TDirectory similar to os.walk.
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.
16 (TDirectory, list(TObject), list(TDirectory))
18 tkeys = tdirectory.GetListOfKeys()
23 tobject = tdirectory.Get(tkey.GetName())
24 if isinstance(tobject, ROOT.TDirectory):
25 tdirectories.append(tobject)
27 tobjects.append(tobject)
29 yield tdirectory, tobjects, tdirectories
32 for sub_tdirectory
in tdirectories:
33 for tdirectory, tobjects, tdirectories
in root_walk(sub_tdirectory):
34 yield tdirectory, tobjects, tdirectories
37 @contextlib.contextmanager
38 def root_open(tfile_or_file_path, tfile_options=None):
39 """Context manager to open a TFile.
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.
46 tfile_or_file_path : str or ROOT.TFile
47 Path to the file or the TFile that should be activated.
49 Option string forwarded to the ROOT.TFile constructor
50 Typical options as "RECREATE", "READ" or "UPDATE".
52 if isinstance(tfile_or_file_path, ROOT.TFile):
53 tfile = tfile_or_file_path
57 save_tdirectory = ROOT.gROOT.CurrentDirectory()
59 if tfile_options
is None:
60 tfile = ROOT.TFile(tfile_or_file_path)
62 tfile = ROOT.TFile(tfile_or_file_path, tfile_options)
71 @contextlib.contextmanager
73 """Context manager that temporarily switches the current global ROOT directory while in the context.
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.
78 If it is not present create it.
82 tdirectory : ROOT.TDirectory or str
83 ROOT directory to switch to or name of a folder to switch.
88 The new current ROOT directory.
94 save_tdirectory = ROOT.gROOT.CurrentDirectory()
96 if not tdirectory
or "." == tdirectory:
97 tdirectory = save_tdirectory
99 elif isinstance(tdirectory, str):
100 tdirectory_name = tdirectory
104 tdirectory = save_tdirectory.GetDirectory(tdirectory_name)
106 tdirectory = save_tdirectory.mkdir(tdirectory_name, tdirectory_name)
108 raise RuntimeError(
"Could not create or find folder %s" % tdirectory_name)
112 tdirectory = save_tdirectory.GetDirectory(tdirectory_name)
115 if tdirectory
is not None:
124 """Strips all meta characters that might be unsafe to use as a ROOT name.
129 A name that should be transformed
134 Name with potentially harmful characters deleted / replaced.
136 deletechars = str.maketrans(
"",
"",
r"/$\#{}()[]=")
137 name = name.replace(
' ',
'_').replace(
'-',
'_').replace(
',',
'_').translate(deletechars)
142 """Open a browser and show the given object.
146 tobject : ROOT.TObject
147 The object to be shown
152 The new TBrowser used to show the object.
158 tbrowser = ROOT.TBrowser()
159 if isinstance(tobject, ROOT.TObject):
160 if hasattr(tobject,
"Browse"):
161 tobject.Browse(tbrowser)
163 tbrowser.BrowseObject(tobject)
165 raise ValueError(
"Can only browse ROOT objects inheriting from TObject.")
171 """Returns a list of names that are contained in the given obj.
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.
180 @root_ls.register(str)
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)
191 @root_ls.register(ROOT.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])
199 @root_ls.register(ROOT.TTree)
200 @root_ls.register(ROOT.TNtuple)
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])