FWLite/python/__init__.py

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