test/unittests/FWLiteAnalyzer.py

0001 #!/usr/bin/env python
0002
0003 # A simple FWLite-based python analyzer
0004 # Based on https://twiki.cern.ch/twiki/bin/view/CMSPublic/WorkBookFWLitePython
0005 # Some snippets of codes are stolen from PhysicsTools/Heppy
0006
0007 from __future__ import print_function
0008 from ROOT import gROOT, gSystem, AutoLibraryLoader
0009 from DataFormats.FWLite import Events, Handle
0010
0011
0012 class FWLiteAnalyzer(object):
0013
0014   def __init__(self, inputFiles=None, handles=None, firstEvent=None, maxEvents=None):
0015     gROOT.SetBatch()        # don't pop up canvases
0016
0017     if inputFiles:
0018       if isinstance(inputFiles, str):
0019         self.inputFiles = [inputFiles]
0020       else:
0021         self.inputFiles = inputFiles
0022     else:
0023       self.inputFiles = []
0024
0025     self.events = Events(self.inputFiles)
0026
0027     self.handles = {}
0028     self.handle_labels = {}
0029     if handles:
0030       for k, v in handles.items():
0031         self.handles[k] = Handle(v[0])
0032         self.handle_labels[k] = v[1]
0033
0034     self.setup = {}
0035
0036     if firstEvent:
0037       self.firstEvent = firstEvent
0038     else:
0039       self.firstEvent = 0
0040
0041     if maxEvents:
0042       self.maxEvents = maxEvents
0043     else:
0044       self.maxEvents = 0x7FFFFFFFFFFFFFFF  # TChain::kBigNumber
0045     return
0046
0047   def analyze(self):
0048     self.beginLoop()
0049     for evt in self.processLoop():
0050       pass
0051     self.endLoop()
0052     return
0053
0054   def beginLoop(self):
0055     return
0056
0057   def endLoop(self):
0058     return
0059
0060   def processLoop(self):
0061     for ievt, evt in enumerate(self.events):
0062       #if (ievt % 1000) == 0:
0063       #  print "Processing event: %i" % ievt
0064       self.process(evt)
0065       yield evt
0066     return
0067
0068   def process(self, event):
0069     self.getHandles(event)
0070     return
0071
0072   def getHandles(self, event):
0073     for k, v in self.handles.items():
0074       label = self.handle_labels[k]
0075       event.getByLabel(label, v)
0076     return
0077
0078
0079 # ______________________________________________________________________________
0080 if __name__ == "__main__":
0081
0082   print("Loading FW Lite")
0083   gSystem.Load("libFWCoreFWLite")
0084   gROOT.ProcessLine("FWLiteEnabler::enable();")
0085
0086   #gSystem.Load("libDataFormatsFWLite.so")
0087   #gSystem.Load("libDataFormatsPatCandidates.so")
0088   #gSystem.Load("libDataFormatsL1TMuon.so")
0089
0090   analyzer = FWLiteAnalyzer(inputFiles='pippo.root')
0091   analyzer.analyze()