PythonAnalysis/python/cmstools.py

0001 """Python helper tools for CMS FWLite
0002
0003 benedikt.hegner@cern.ch
0004
0005 """
0006 from builtins import range
0007 import re
0008 import ROOT
0009 import sys
0010 ### define tab completion
0011 try:
0012   import readline #cmscompleter
0013   readline.parse_and_bind('tab: complete')
0014 except:
0015   print('WARNING: Could not load tab completion')
0016
0017
0018 # for adding iterators at runtime
0019 from . import iterators
0020
0021
0022 ### workaround iterator generators for ROOT classes
0023 def all(container):
0024
0025   # loop over ROOT::TTree and similar
0026   if hasattr(container,'GetEntries'):
0027     try:
0028       entries = container.GetEntries()
0029       for entry in range(entries):
0030         yield entry
0031     except:
0032         raise cmserror("Looping of %s failed" %container)
0033
0034   # loop over std::vectors and similar
0035   elif hasattr(container, 'size'):
0036     try:
0037       entries = container.size()
0038       for entry in range(entries):
0039         yield container[entry]
0040     except:
0041       pass
0042
0043  # loop over containers with begin and end iterators
0044 def loop(begin, end):
0045     """Convert a pair of C++ iterators into a python generator"""
0046     while (begin != end):
0047         yield begin.__deref__()  #*b
0048         begin.__preinc__()       #++b
0049
0050 ### auto branch types (Chris Jones)
0051 def createBranchBuffer(branch):
0052     reColons = re.compile(r'::')
0053     reCloseTemplate =re.compile(r'>')
0054     reOpenTemplate =re.compile(r'<')
0055     branchType = ROOT.branchToClass(branch)
0056     #buffer = eval ('ROOT.'+reColons.sub(".",reOpenTemplate.sub("(ROOT.",reCloseTemplate.sub(")",branchType.GetName())))+'()')
0057     buffer = ROOT.MakeRootClass(branchType.GetName()) ()
0058     if( branch.GetName()[-1] != '.') and (branch.GetName()!="EventAuxiliary"):
0059         branch.SetAddress(buffer)
0060     else:
0061         branch.SetAddress(ROOT.AddressOf(buffer))
0062     return buffer
0063
0064
0065 class EventTree(object):
0066       def __init__(self,obj):
0067           sys.stderr.write ("WARNING: This package has been deprecated and will be removed in the near future.\nPlease switch to using FWLite.Python (https://twiki.cern.ch/twiki/bin/viewauth/CMS/WorkBookFWLitePython)\n")
0068           treeName = 'Events'
0069           if isinstance(obj, ROOT.TTree):
0070               self._tree = obj
0071           elif isinstance(obj, ROOT.TFile):
0072               self._tree = obj.Get(treeName)
0073           elif isinstance(obj, str):
0074               self._tree = ROOT.TFile.Open(obj).Get(treeName)
0075           else:
0076               raise cmserror("EventTree accepts only TTrees, TFiles and filenames")
0077           self._usedBranches = dict()
0078           self._index = -1
0079           self._aliases = self._tree.GetListOfAliases()
0080       def branch(self,name):
0081           # support for aliases
0082           alias = self._tree.GetAlias(name)
0083           if alias != '': name = alias
0084           # access the branch in ttree
0085           if name in self._usedBranches:
0086               return self._usedBranches[name]
0087           self._usedBranches[name]=EventBranch(self,name)
0088           return self._usedBranches[name]
0089       def cppCode(self, name):
0090           """C++ code for accessing the product inside the full framework"""
0091           alias = self._tree.GetAlias(name)
0092           if alias != '': name = alias
0093           tmpBranch = self._tree.GetBranch(name)
0094           typeString = ROOT.branchToClass(tmpBranch).GetName()
0095           if "edm::Wrapper" in typeString:
0096               typeString = typeString.replace("<edm::Wrapper","")
0097               typeString = typeString.rstrip(">")
0098           nameParts = name.split("_")
0099           if nameParts[2] == "":
0100               cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", dummy);'\
0101                         %(typeString, nameParts[1])
0102           else:
0103               cppCode = 'edm::Handle<%s > dummy;\nevent.getByLabel("%s", "%s", dummy);'\
0104                         %(typeString, nameParts[1], nameParts[2])
0105           return cppCode
0106       def getListOfAliases(self):
0107           return self._aliases
0108       def SetAlias (self, alias, fullName):
0109           self.tree().SetAlias(alias, fullName)
0110       def index(self):
0111           return self._index
0112       def tree(self):
0113           return self._tree
0114       def __setBranchIndicies(self):
0115           for branch in self._usedBranches.values():
0116               branch.setIndex(self._index)
0117       def __getattr__(self, name):
0118           return self.branch(name)
0119       def __getitem__(self,key):
0120           if key <0 or key > self._tree.GetEntries():
0121               raise IndexError
0122           self._index = key
0123           self.__setBranchIndicies()
0124           self._tree.GetEntry(self._index,0)
0125           return Event(self)
0126       def __iter__(self):
0127           # flushing/initializing the root buffers
0128           entry = 0
0129           self._index = entry
0130           self.__setBranchIndicies()
0131           self._tree.GetEntry(self._index,0)
0132           # the real loop
0133           for entry in range(self._tree.GetEntries()):
0134               self._index = entry
0135               self.__setBranchIndicies()
0136               self._tree.GetEntry(self._index,0)
0137               yield Event(self)    # TODO: don't return a new object but update the old one
0138
0139
0140 class Event(object):
0141     def __init__(self, eventTree):
0142         self._eventTree = eventTree
0143
0144     def getProduct(self, name):
0145         return iterators.addIterator(self._eventTree.branch(name)())
0146
0147     def __getattr__(self, name):
0148         return iterators.addIterator(self._eventTree.branch(name)())
0149
0150
0151 class EventBranch(object):
0152     def __init__(self,parent,name):
0153         self._branch = parent.tree().GetBranch(name)
0154         if self._branch == None:
0155             raise cmserror("Unknown branch "+name)
0156         self._buffer = createBranchBuffer(self._branch)
0157         self._index = parent.index()
0158         self._readData = False
0159     def setIndex(self,index):
0160         self._index = index
0161         self._readData = False
0162     def __readData(self):
0163         self._branch.GetEntry(self._index)
0164         self._readData = True
0165
0166     # replace this by __getattr__ to allow branch.attr instead of branch().attr
0167     def __call__(self):
0168         if not self._readData:
0169             self.__readData()
0170         return self._buffer
0171
0172
0173 class cmserror(Exception):
0174     def __init__(self, message):
0175           length = len(message)+7   #7=len("ERROR: ")
0176           print("="*length)
0177           print("ERROR:", message)
0178           print("="*length)