PythonAnalysis/python/cmstools.py

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