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