Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-10-25 09:58:08

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  # don't initialize twice
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             # force calling SetEntry as a new ValueReader was created
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             # force calling SetEntry as a new ArrayReader was created
0094             tree.gotoEntry(tree.entry, forceCall=True)
0095             return _ar
0096 
0097 
0098 ####### PRIVATE IMPLEMENTATION PART #######
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