File indexing completed on 2024-04-06 12:23:46
0001 import types
0002 import ROOT
0003 ROOT.PyConfig.IgnoreCommandLineOptions = True
0004
0005
0006 def InputTree(tree, entrylist=ROOT.MakeNullPointer(ROOT.TEntryList)):
0007 """add to the PyROOT wrapper of a TTree a TTreeReader and methods readBranch, arrayReader, valueReader"""
0008 if hasattr(tree, '_ttreereader'):
0009 return tree
0010 tree.entry = -1
0011 tree._entrylist = entrylist
0012 tree._ttreereader = ROOT.TTreeReader(tree, tree._entrylist)
0013 tree._ttreereader._isClean = True
0014 tree._ttrvs = {}
0015 tree._ttras = {}
0016 tree._leafTypes = {}
0017 tree._ttreereaderversion = 1
0018 tree.arrayReader = types.MethodType(getArrayReader, tree)
0019 tree.valueReader = types.MethodType(getValueReader, tree)
0020 tree.readBranch = types.MethodType(readBranch, tree)
0021 tree.gotoEntry = types.MethodType(_gotoEntry, tree)
0022 tree.readAllBranches = types.MethodType(_readAllBranches, tree)
0023 tree.entries = tree._ttreereader.GetEntries(False)
0024 tree._extrabranches = {}
0025 return tree
0026
0027
0028 def getArrayReader(tree, branchName):
0029 """Make a reader for branch branchName containing a variable-length value array."""
0030 if branchName not in tree._ttras:
0031 if not tree.GetBranch(branchName):
0032 raise RuntimeError("Can't find branch '%s'" % branchName)
0033 if not tree.GetBranchStatus(branchName):
0034 raise RuntimeError("Branch %s has status=0" % branchName)
0035 leaf = tree.GetBranch(branchName).GetLeaf(branchName)
0036 if not bool(leaf.GetLeafCount()):
0037 raise RuntimeError("Branch %s is not a variable-length value array" % branchName)
0038 typ = leaf.GetTypeName()
0039 tree._ttras[branchName] = _makeArrayReader(tree, typ, branchName)
0040 return tree._ttras[branchName]
0041
0042
0043 def getValueReader(tree, branchName):
0044 """Make a reader for branch branchName containing a single value."""
0045 if branchName not in tree._ttrvs:
0046 if not tree.GetBranch(branchName):
0047 raise RuntimeError("Can't find branch '%s'" % branchName)
0048 if not tree.GetBranchStatus(branchName):
0049 raise RuntimeError("Branch %s has status=0" % branchName)
0050 leaf = tree.GetBranch(branchName).GetLeaf(branchName)
0051 if bool(leaf.GetLeafCount()) or leaf.GetLen() != 1:
0052 raise RuntimeError("Branch %s is not a value" % branchName)
0053 typ = leaf.GetTypeName()
0054 tree._ttrvs[branchName] = _makeValueReader(tree, typ, branchName)
0055 return tree._ttrvs[branchName]
0056
0057
0058 def clearExtraBranches(tree):
0059 tree._extrabranches = {}
0060
0061
0062 def setExtraBranch(tree, name, val):
0063 tree._extrabranches[name] = val
0064
0065
0066 def readBranch(tree, branchName):
0067 """Return the branch value if the branch is a value, and a TreeReaderArray if the branch is an array"""
0068 if tree._ttreereader._isClean:
0069 raise RuntimeError("readBranch must not be called before calling gotoEntry")
0070 if branchName in tree._extrabranches:
0071 return tree._extrabranches[branchName]
0072 elif branchName in tree._ttras:
0073 return tree._ttras[branchName]
0074 elif branchName in tree._ttrvs:
0075 ret = tree._ttrvs[branchName].Get()[0]
0076 return ord(ret) if type(ret) == str else ret
0077 else:
0078 branch = tree.GetBranch(branchName)
0079 if not branch:
0080 raise RuntimeError("Unknown branch %s" % branchName)
0081 if not tree.GetBranchStatus(branchName):
0082 raise RuntimeError("Branch %s has status=0" % branchName)
0083 leaf = branch.GetLeaf(branchName)
0084 typ = leaf.GetTypeName()
0085 if leaf.GetLen() == 1 and not bool(leaf.GetLeafCount()):
0086 _vr = _makeValueReader(tree, typ, branchName)
0087
0088 tree.gotoEntry(tree.entry, forceCall=True)
0089 ret = _vr.Get()[0]
0090 return ord(ret) if type(ret) == str else ret
0091 else:
0092 _ar = _makeArrayReader(tree, typ, branchName)
0093
0094 tree.gotoEntry(tree.entry, forceCall=True)
0095 return _ar
0096
0097
0098
0099
0100 def _makeArrayReader(tree, typ, nam):
0101 if not tree._ttreereader._isClean:
0102 _remakeAllReaders(tree)
0103 ttra = ROOT.TTreeReaderArray(typ)(tree._ttreereader, nam)
0104 tree._leafTypes[nam] = typ
0105 tree._ttras[nam] = ttra
0106 return tree._ttras[nam]
0107
0108
0109 def _makeValueReader(tree, typ, nam):
0110 if not tree._ttreereader._isClean:
0111 _remakeAllReaders(tree)
0112 ttrv = ROOT.TTreeReaderValue(typ)(tree._ttreereader, nam)
0113 tree._leafTypes[nam] = typ
0114 tree._ttrvs[nam] = ttrv
0115 return tree._ttrvs[nam]
0116
0117
0118 def _remakeAllReaders(tree):
0119 _ttreereader = ROOT.TTreeReader(tree, getattr(tree, '_entrylist', ROOT.MakeNullPointer(ROOT.TEntryList)))
0120 _ttreereader._isClean = True
0121 _ttrvs = {}
0122 for k in tree._ttrvs.keys():
0123 _ttrvs[k] = ROOT.TTreeReaderValue(tree._leafTypes[k])(_ttreereader, k)
0124 _ttras = {}
0125 for k in tree._ttras.keys():
0126 _ttras[k] = ROOT.TTreeReaderArray(tree._leafTypes[k])(_ttreereader, k)
0127 tree._ttrvs = _ttrvs
0128 tree._ttras = _ttras
0129 tree._ttreereader = _ttreereader
0130 tree._ttreereaderversion += 1
0131
0132
0133 def _readAllBranches(tree):
0134 tree.GetEntry(_currentTreeEntry(tree))
0135
0136
0137 def _currentTreeEntry(tree):
0138 if tree._entrylist:
0139 return tree._entrylist.GetEntry(tree.entry)
0140 else:
0141 return tree.entry
0142
0143
0144 def _gotoEntry(tree, entry, forceCall=False):
0145 tree._ttreereader._isClean = False
0146 if tree.entry != entry or forceCall:
0147 if (tree.entry == entry - 1 and entry != 0):
0148 tree._ttreereader.Next()
0149 else:
0150 tree._ttreereader.SetEntry(entry)
0151 tree.entry = entry
