Back to home page

Project CMSSW displayed by LXR

 
 

    


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