FWLite/python/__init__.py

0001 #! /usr/bin/env python
0002
0003 from __future__ import print_function
0004 import ROOT
0005 import inspect
0006 import sys
0007 import optparse
0008 from FWCore.ParameterSet.VarParsing import VarParsing
0009 from builtins import int
0010
0011 ROOT.gSystem.Load("libFWCoreFWLite")
0012 ROOT.FWLiteEnabler.enable()
0013
0014 # Whether warn() should print anythingg
0015 quietWarn = False
0016
0017 def setQuietWarn (quiet = True):
0018     global quietWarn
0019     quietWarn = quiet
0020
0021 def warn (*args, **kwargs):
0022     """print out warning with line number and rest of arguments"""
0023     if quietWarn: return
0024     frame = inspect.stack()[1]
0025     filename = frame[1]
0026     lineNum  = frame[2]
0027     #print "after '%s'" % filename
0028     blankLines = kwargs.get('blankLines', 0)
0029     if blankLines:
0030         print('\n' * blankLines)
0031     spaces = kwargs.get('spaces', 0)
0032     if spaces:
0033         print(' ' * spaces, end=' ')
0034     if len (args):
0035         print("%s (%s): " % (filename, lineNum), end=' ')
0036         for arg in args:
0037             print(arg, end=' ')
0038         print()
0039     else:
0040         print("%s (%s):" % (filename, lineNum))
0041
0042 ########################
0043 ## ################## ##
0044 ## ## ############ ## ##
0045 ## ## ## Handle ## ## ##
0046 ## ## ############ ## ##
0047 ## ################## ##
0048 ########################
0049
0050 class Handle:
0051     """Python interface to FWLite Handle class"""
0052
0053     def __init__ (self,
0054                   typeString,
0055                   **kwargs):
0056         """Initialize python handle wrapper """
0057         # turn off warnings
0058         oldWarningLevel = ROOT.gErrorIgnoreLevel
0059         ROOT.gErrorIgnoreLevel = ROOT.kError
0060         self._nodel = False
0061         if kwargs.get ('noDelete'):
0062             print("Not deleting wrapper")
0063             del kwargs['noDelete']
0064         else:
0065             self._nodel = True
0066         self._type = typeString
0067         self._resetWrapper()
0068         self._exception = RuntimeError ("getByLabel not called for '%s'", self)
0069         # restore warning state
0070         ROOT.gErrorIgnoreLevel = oldWarningLevel
0071         # Since we deleted the options as we used them, that means
0072         # that kwargs should be empty.  If it's not, that means that
0073         # somebody passed in an argument that we're not using and we
0074         # should complain.
0075         if len (kwargs):
0076             raise RuntimeError("Unknown arguments %s" % kwargs)
0077
0078     def isValid (self):
0079         """Returns true if getByLabel call was successful and data is
0080         present in handle."""
0081         return not self._exception
0082
0083
0084     def product (self):
0085         """Returns product stored in handle."""
0086         if self._exception:
0087             raise self._exception
0088         return self._wrapper.product()
0089
0090
0091     def __str__ (self):
0092         return "%s" % (self._type)
0093
0094
0095     ## Private member functions ##
0096
0097     def _resetWrapper (self):
0098         """(Internal) reset the edm wrapper"""
0099         self._wrapper   = ROOT.edm.Wrapper (self._type)()
0100         self._typeInfo  = self._wrapper.typeInfo()
0101         ROOT.SetOwnership (self._wrapper, False)
0102         # O.k.  This is a little weird.  We want a pointer to an EDM
0103         # wrapper, but we don't want the memory it is pointing to.
0104         # So, we've created it and grabbed the type info.  Since we
0105         # don't want a memory leak, we destroy it.
0106         if not self._nodel :
0107             ROOT.TClass.GetClass("edm::Wrapper<"+self._type+">").Destructor( self._wrapper )
0108
0109     def _typeInfoGetter (self):
0110         """(Internal) Return the type info"""
0111         return self._typeInfo
0112
0113
0114     def _addressOf (self):
0115         """(Internal) Return address of edm wrapper"""
0116         return ROOT.AddressOf (self._wrapper)
0117
0118
0119     def _setStatus (self, getByLabelSuccess, labelString):
0120         """(Internal) To be called by Events.getByLabel"""
0121         if not getByLabelSuccess:
0122             self._exception = RuntimeError ("getByLabel (%s, %s) failed" \
0123                                             % (self, labelString))
0124             return
0125         if not self._wrapper.isPresent():
0126             self._exception = RuntimeError ("getByLabel (%s, %s) not present this event" \
0127                                             % (self, labelString))
0128             return
0129         # if we're still here, then everything is happy.  Clear the exception
0130         self._exception = None
0131
0132
0133 #######################
0134 ## ################# ##
0135 ## ## ########### ## ##
0136 ## ## ## Lumis ## ## ##
0137 ## ## ########### ## ##
0138 ## ################# ##
0139 #######################
0140
0141 class Lumis:
0142     """Python interface to FWLite LuminosityBlock"""
0143     def __init__ (self, inputFiles = '', **kwargs):
0144         self._lumi = None
0145         self._lumiCounts = 0
0146         self._tfile = None
0147         self._maxLumis = 0
0148         if isinstance (inputFiles, list):
0149             # it's a list
0150             self._filenames = inputFiles[:]
0151         elif isinstance (inputFiles, VarParsing):
0152             # it's a VarParsing object
0153             options = inputFiles
0154             self._maxLumis           = options.maxEvents
0155             self._filenames          = options.inputFiles
0156         else:
0157             # it's probably a single string
0158             self._filenames = [inputFiles]
0159         ##############################
0160         ## Parse optional arguments ##
0161         ##############################
0162         if 'maxEvents' in kwargs:
0163             self._maxLumis = kwargs['maxEvents']
0164             del kwargs['maxEvents']
0165         if 'options' in kwargs:
0166             options = kwargs ['options']
0167             self._maxLumis           = options.maxEvents
0168             self._filenames          = options.inputFiles
0169             self._secondaryFilenames = options.secondaryInputFiles
0170             del kwargs['options']
0171         # Since we deleted the options as we used them, that means
0172         # that kwargs should be empty.  If it's not, that means that
0173         # somebody passed in an argument that we're not using and we
0174         # should complain.
0175         if len (kwargs):
0176             raise RuntimeError("Unknown arguments %s" % kwargs)
0177         if not self._filenames:
0178             raise RuntimeError("No input files given")
0179         if not self._createFWLiteLumi():
0180             # this shouldn't happen as you are getting nothing the
0181             # very first time out, but let's at least check to
0182             # avoid problems.
0183             raise RuntimeError("Never and information about Lumi")
0184
0185
0186     def __del__ (self):
0187         """(Internal) Destructor"""
0188         # print "Goodbye cruel world, I'm leaving you today."
0189         del self._lumi
0190         # print "Goodbye, goodbye, goodbye."
0191
0192
0193     def __iter__ (self):
0194         return self._next()
0195
0196
0197     def aux (self):
0198         try:
0199             return self._lumi.luminosityBlockAuxiliary()
0200         except:
0201             raise RuntimeError("Lumis.aux() called on object in invalid state")
0202
0203
0204     def luminosityBlockAuxiliary (self):
0205         try:
0206             return self._lumi.luminosityBlockAuxiliary()
0207         except:
0208             raise RuntimeError("Lumis.luminosityBlockAuxiliary() called on object in invalid state")
0209
0210
0211     def getByLabel (self, *args):
0212         """Calls FWLite's getByLabel.  Called:
0213         getByLabel (moduleLabel, handle)
0214         getByLabel (moduleLabel, productInstanceLabel, handle),
0215         getByLabel (moduleLabel, productInstanceLabel, processLabel, handle),
0216         or
0217         getByLabel ( (mL, pIL,pL), handle)
0218         """
0219         length = len (args)
0220         if length < 2 or length > 4:
0221             # not called correctly
0222             raise RuntimeError("Incorrect number of arguments")
0223         # handle is always the last argument
0224         argsList = list (args)
0225         handle = argsList.pop()
0226         if len(argsList)==1 :
0227             if( isinstance (argsList[0], tuple) or
0228                 isinstance (argsList[0], list) ) :
0229                 if len (argsList[0]) > 3:
0230                     raise RuntimeError("getByLabel Error: label tuple has too " \
0231                         "many arguments '%s'" % argsList[0])
0232                 argsList = list(argsList[0])
0233             if( isinstance(argsList[0], str) and ":" in argsList[0] ):
0234                 if argsList[0].count(":") > 3:
0235                     raise RuntimeError("getByLabel Error: label tuple has too " \
0236                         "many arguments '%s'" % argsList[0].split(":"))
0237                 argsList = argsList[0].split(":")
0238         while len(argsList) < 3:
0239             argsList.append ('')
0240         (moduleLabel, productInstanceLabel, processLabel) = argsList
0241         labelString = "'" + "', '".join(argsList) + "'"
0242         if not handle._wrapper :
0243             handle._resetWrapper()
0244         handle._setStatus ( self._lumi.getByLabel( handle._typeInfoGetter(),
0245                                                    moduleLabel,
0246                                                    productInstanceLabel,
0247                                                    processLabel,
0248                                                    handle._addressOf() ),
0249                             labelString )
0250         return handle.isValid()
0251
0252
0253     ##############################
0254     ## Private Member Functions ##
0255     ##############################
0256
0257     def _createFWLiteLumi (self):
0258         """(Internal) Creates an FWLite Lumi"""
0259         # are there any files left?
0260         if not self._filenames:
0261             return False
0262         if self._lumi:
0263             del self._lumi
0264             self._lumi = None
0265         self._veryFirstTime = False
0266         self._currFilename = self._filenames.pop(0)
0267         #print "Opening file", self._currFilename
0268         if self._tfile:
0269             del self._tfile
0270         self._tfile = ROOT.TFile.Open (self._currFilename)
0271         self._lumi = ROOT.fwlite.LuminosityBlock (self._tfile);
0272         self._lumi.toBegin()
0273         return True
0274
0275
0276     def _next (self):
0277         """(Internal) Iterator internals"""
0278         while True:
0279             if self._lumi.atEnd():
0280                 if not self._createFWLiteLumi():
0281                     # there are no more files here, so we are done
0282                     break
0283             yield self
0284             self._lumiCounts += 1
0285             if self._maxLumis > 0 and self._lumiCounts >= self._maxLumis:
0286                 break
0287             self._lumi.__preinc__()
0288
0289
0290
0291 ######################
0292 ## ################ ##
0293 ## ## ########## ## ##
0294 ## ## ## Runs ## ## ##
0295 ## ## ########## ## ##
0296 ## ################ ##
0297 ######################
0298
0299 class Runs:
0300     """Python interface to FWLite LuminosityBlock"""
0301     def __init__ (self, inputFiles = '', **kwargs):
0302         self._run = None
0303         self._runCounts = 0
0304         self._tfile = None
0305         self._maxRuns = 0
0306         if isinstance (inputFiles, list):
0307             # it's a list
0308             self._filenames = inputFiles[:]
0309         elif isinstance (inputFiles, VarParsing):
0310             # it's a VarParsing object
0311             options = inputFiles
0312             self._maxRuns           = options.maxEvents
0313             self._filenames           = options.inputFiles
0314         else:
0315             # it's probably a single string
0316             self._filenames = [inputFiles]
0317         ##############################
0318         ## Parse optional arguments ##
0319         ##############################
0320         if 'maxEvents' in kwargs:
0321             self._maxRuns = kwargs['maxEvents']
0322             del kwargs['maxEvents']
0323         if 'options' in kwargs:
0324             options = kwargs ['options']
0325             self._maxRuns           = options.maxEvents
0326             self._filenames           = options.inputFiles
0327             self._secondaryFilenames  = options.secondaryInputFiles
0328             del kwargs['options']
0329         # Since we deleted the options as we used them, that means
0330         # that kwargs should be empty.  If it's not, that means that
0331         # somebody passed in an argument that we're not using and we
0332         # should complain.
0333         if len (kwargs):
0334             raise RuntimeError("Unknown arguments %s" % kwargs)
0335         if not self._filenames:
0336             raise RuntimeError("No input files given")
0337         if not self._createFWLiteRun():
0338             # this shouldn't happen as you are getting nothing the
0339             # very first time out, but let's at least check to
0340             # avoid problems.
0341             raise RuntimeError("Never and information about Run")
0342
0343
0344     def __del__ (self):
0345         """(Internal) Destructor"""
0346         # print "Goodbye cruel world, I'm leaving you today."
0347         del self._run
0348         # print "Goodbye, goodbye, goodbye."
0349
0350
0351     def __iter__ (self):
0352         return self._next()
0353
0354
0355     def aux (self):
0356         try:
0357             return self._run.runAuxiliary()
0358         except:
0359             raise RuntimeError("Runs.aux() called on object in invalid state")
0360
0361
0362     def runAuxiliary (self):
0363         try:
0364             return self._run.runAuxiliary()
0365         except:
0366             raise RuntimeError("Runs.runAuxiliary() called on object in invalid state")
0367
0368
0369     def getByLabel (self, *args):
0370         """Calls FWLite's getByLabel.  Called:
0371         getByLabel (moduleLabel, handle)
0372         getByLabel (moduleLabel, productInstanceLabel, handle),
0373         getByLabel (moduleLabel, productInstanceLabel, processLabel, handle),
0374         or
0375         getByLabel ( (mL, pIL,pL), handle)
0376         """
0377         length = len (args)
0378         if length < 2 or length > 4:
0379             # not called correctly
0380             raise RuntimeError("Incorrect number of arguments")
0381         # handle is always the last argument
0382         argsList = list (args)
0383         handle = argsList.pop()
0384         if len(argsList)==1 :
0385             if( isinstance (argsList[0], tuple) or
0386                 isinstance (argsList[0], list) ) :
0387                 if len (argsList[0]) > 3:
0388                     raise RuntimeError("getByLabel Error: label tuple has too " \
0389                         "many arguments '%s'" % argsList[0])
0390                 argsList = list(argsList[0])
0391             if( isinstance(argsList[0], str) and ":" in argsList[0] ):
0392                 if argsList[0].count(":") > 3:
0393                     raise RuntimeError("getByLabel Error: label tuple has too " \
0394                         "many arguments '%s'" % argsList[0].split(":"))
0395                 argsList = argsList[0].split(":")
0396         while len(argsList) < 3:
0397             argsList.append ('')
0398         (moduleLabel, productInstanceLabel, processLabel) = argsList
0399         labelString = "'" + "', '".join(argsList) + "'"
0400         if not handle._wrapper :
0401             handle._resetWrapper()
0402         handle._setStatus ( self._run.getByLabel( handle._typeInfoGetter(),
0403                                                    moduleLabel,
0404                                                    productInstanceLabel,
0405                                                    processLabel,
0406                                                    handle._addressOf() ),
0407                             labelString )
0408         return handle.isValid()
0409
0410
0411
0412
0413     ##############################
0414     ## Private Member Functions ##
0415     ##############################
0416
0417     def _createFWLiteRun (self):
0418         """(Internal) Creates an FWLite Run"""
0419         # are there any files left?
0420         if not self._filenames:
0421             return False
0422         if self._run:
0423             del self._run
0424             self._run = None
0425         self._veryFirstTime = False
0426         self._currFilename = self._filenames.pop(0)
0427         #print "Opening file", self._currFilename
0428         if self._tfile:
0429             del self._tfile
0430         self._tfile = ROOT.TFile.Open (self._currFilename)
0431         self._run = ROOT.fwlite.Run (self._tfile);
0432         self._run.toBegin()
0433         return True
0434
0435
0436     def _next (self):
0437         """(Internal) Iterator internals"""
0438         while True:
0439             if self._run.atEnd():
0440                 if not self._createFWLiteRun():
0441                     # there are no more files here, so we are done
0442                     break
0443             yield self
0444             self._runCounts += 1
0445             if self._maxRuns > 0 and self._runCounts >= self._maxRuns:
0446                 break
0447             self._run.__preinc__()
0448
0449
0450 ########################
0451 ## ################## ##
0452 ## ## ############ ## ##
0453 ## ## ## Events ## ## ##
0454 ## ## ############ ## ##
0455 ## ################## ##
0456 ########################
0457
0458 class Events:
0459     """Python interface to FWLite ChainEvent class"""
0460
0461     def __init__(self, inputFiles = '', **kwargs):
0462         """inputFiles    => Either a single filename or a list of filenames
0463         Optional arguments:
0464         forceEvent  => Use fwlite::Event IF there is only one file
0465         maxEvents   => Maximum number of events to process
0466         """
0467         self._veryFirstTime      = True
0468         self._event              = 0
0469         self._eventCounts        = 0
0470         self._maxEvents          = 0
0471         self._forceEvent         = False
0472         self._mode               = None
0473         self._secondaryFilenames = None
0474         if isinstance (inputFiles, list):
0475             # it's a list
0476             self._filenames = inputFiles[:]
0477         elif isinstance (inputFiles, VarParsing):
0478             # it's a VarParsing object
0479             options = inputFiles
0480             self._maxEvents           = options.maxEvents
0481             self._filenames           = options.inputFiles
0482             self._secondaryFilenames  = options.secondaryInputFiles
0483         else:
0484             # it's probably a single string
0485             self._filenames = [inputFiles]
0486         ##############################
0487         ## Parse optional arguments ##
0488         ##############################
0489         if 'maxEvents' in kwargs:
0490             self._maxEvents = kwargs['maxEvents']
0491             del kwargs['maxEvents']
0492         if 'forceEvent' in kwargs:
0493             self._forceEvent = kwargs['forceEvent']
0494             del kwargs['forceEvent']
0495         if 'options' in kwargs:
0496             options = kwargs ['options']
0497             self._maxEvents           = options.maxEvents
0498             self._filenames           = options.inputFiles
0499             self._secondaryFilenames  = options.secondaryInputFiles
0500             del kwargs['options']
0501         # Since we deleted the options as we used them, that means
0502         # that kwargs should be empty.  If it's not, that means that
0503         # somebody passed in an argument that we're not using and we
0504         # should complain.
0505         if len (kwargs):
0506             raise RuntimeError("Unknown arguments %s" % kwargs)
0507         if not self._filenames:
0508             raise RuntimeError("No input files given")
0509
0510
0511     def to (self, entryIndex):
0512         """Jumps to event entryIndex"""
0513         if self._veryFirstTime:
0514             self._createFWLiteEvent()
0515         return self._event.to ( int(entryIndex) )
0516
0517
0518     def toBegin (self):
0519         """Called to reset event loop to first event."""
0520         self._toBegin = True
0521
0522
0523     def size (self):
0524         """Returns number of events"""
0525         if self._veryFirstTime:
0526             self._createFWLiteEvent()
0527         return self._event.size()
0528
0529
0530     def eventAuxiliary (self):
0531         """Returns eventAuxiliary object"""
0532         if self._veryFirstTime:
0533             raise RuntimeError("eventAuxiliary() called before "\
0534                   "toBegin() or to()")
0535         return self._event.eventAuxiliary()
0536
0537
0538     def object (self):
0539         """Returns event object"""
0540         return self._event
0541
0542
0543     def getByLabel (self, *args):
0544         """Calls FWLite's getByLabel.  Called:
0545         getByLabel (moduleLabel, handle)
0546         getByLabel (moduleLabel, productInstanceLabel, handle),
0547         getByLabel (moduleLabel, productInstanceLabel, processLabel, handle),
0548         or
0549         getByLabel ( (mL, pIL,pL), handle)
0550         """
0551         if self._veryFirstTime:
0552             self._createFWLiteEvent()
0553         length = len (args)
0554         if length < 2 or length > 4:
0555             # not called correctly
0556             raise RuntimeError("Incorrect number of arguments")
0557         # handle is always the last argument
0558         argsList = list (args)
0559         handle = argsList.pop()
0560         if len(argsList)==1 :
0561             if( isinstance (argsList[0], tuple) or
0562                 isinstance (argsList[0], list) ) :
0563                 if len (argsList[0]) > 3:
0564                     raise RuntimeError("getByLabel Error: label tuple has too " \
0565                         "many arguments '%s'" % argsList[0])
0566                 argsList = list(argsList[0])
0567             if( isinstance(argsList[0], str) and ":" in argsList[0] ):
0568                 if argsList[0].count(":") > 3:
0569                     raise RuntimeError("getByLabel Error: label tuple has too " \
0570                         "many arguments '%s'" % argsList[0].split(":"))
0571                 argsList = argsList[0].split(":")
0572         while len(argsList) < 3:
0573             argsList.append ('')
0574         (moduleLabel, productInstanceLabel, processLabel) = argsList
0575         labelString = "'" + "', '".join(argsList) + "'"
0576         if not handle._wrapper :
0577             handle._resetWrapper()
0578         handle._setStatus ( self._event.getByLabel( handle._typeInfoGetter(),
0579                                                     moduleLabel,
0580                                                     productInstanceLabel,
0581                                                     processLabel,
0582                                                     handle._addressOf() ),
0583                             labelString )
0584         return handle.isValid()
0585
0586
0587     def __iter__ (self):
0588         return self._next()
0589
0590
0591     def fileIndex (self):
0592         if self._event:
0593             return self._event.fileIndex()
0594         else:
0595             # default non-existant value is -1.  Return something else
0596             return -2
0597
0598
0599     def secondaryFileIndex (self):
0600         if self._event:
0601             return self._event.secondaryFileIndex()
0602         else:
0603             # default non-existant value is -1.  Return something else
0604             return -2
0605
0606
0607     def fileIndicies (self):
0608         return (self.fileIndex(), self.secondaryFileIndex())
0609
0610
0611     ## Private Member Functions ##
0612
0613
0614     def _parseOptions (self, options):
0615         """(Internal) Parse options"""
0616
0617
0618     def _toBeginCode (self):
0619         """(Internal) Does actual work of toBegin() call"""
0620         self._toBegin = False
0621         self._event.toBegin()
0622         self._eventCounts = 0
0623
0624
0625     def __del__ (self):
0626         """(Internal) Destructor"""
0627         # print "Goodbye cruel world, I'm leaving you today."
0628         del self._event
0629         # print "Goodbye, goodbye, goodbye."
0630
0631
0632     def _createFWLiteEvent (self):
0633         """(Internal) Creates an FWLite Event"""
0634         self._veryFirstTime = False
0635         self._toBegin = True
0636         if isinstance (self._filenames[0], ROOT.TFile):
0637             self._event = ROOT.fwlite.Event (self._filenames[0])
0638             self._mode = 'single'
0639             return self._mode
0640         if len (self._filenames) == 1 and self._forceEvent:
0641             self._tfile = ROOT.TFile.Open (self._filenames[0])
0642             self._event = ROOT.fwlite.Event (self._tfile)
0643             self._mode = 'single'
0644             return self._mode
0645         filenamesSVec = ROOT.vector("string") ()
0646         for name in self._filenames:
0647             filenamesSVec.push_back (name)
0648         if self._secondaryFilenames:
0649             secondarySVec =  ROOT.vector("string") ()
0650             for name in self._secondaryFilenames:
0651                 secondarySVec.push_back (name)
0652             self._event = ROOT.fwlite.MultiChainEvent (filenamesSVec,
0653                                                        secondarySVec)
0654             self._mode = 'multi'
0655         else:
0656             self._event = ROOT.fwlite.ChainEvent (filenamesSVec)
0657             self._mode = 'chain'
0658         return self._mode
0659
0660
0661     def _next (self):
0662         """(Internal) Iterator internals"""
0663         if self._veryFirstTime:
0664             self._createFWLiteEvent()
0665         if self._toBegin:
0666             self._toBeginCode()
0667         while not self._event.atEnd() :
0668             yield self
0669             self._eventCounts += 1
0670             if self._maxEvents > 0 and self._eventCounts >= self._maxEvents:
0671                 break
0672             # Have we been asked to go to the first event?
0673             if self._toBegin:
0674                 self._toBeginCode()
0675             else:
0676                 # if not, lets go to the next event
0677                 self._event.__preinc__()
0678
0679
0680
0681 if __name__ == "__main__":
0682     # test code can go here
0683     pass