ParameterSet/python/VarParsing.py

0001 from __future__ import print_function
0002 import sys
0003 import os
0004 import re
0005 from pprint import pprint
0006 from FWCore.Utilities.Enumerate import Enumerate
0007 from FWCore.Utilities.FileUtils import sectionNofTotal
0008
0009 class VarParsing (object):
0010     """Infrastructure to parse variable definitions passed to cmsRun
0011     configuration scripts"""
0012
0013
0014     multiplicity = Enumerate ("singleton list", "multiplicity")
0015     varType      = Enumerate ("bool int float string tagString")
0016     commaRE      = re.compile (r',')
0017     trueRE       = re.compile (r'^true$',  re.IGNORECASE)
0018     falseRE      = re.compile (r'^false$', re.IGNORECASE)
0019
0020
0021     def __init__ (self, *args):
0022         """Class initializer"""
0023         # Set everything up first
0024         self._singletons       = {}
0025         self._lists            = {}
0026         self._register         = {}
0027         self._beenSet          = {}
0028         self._info             = {}
0029         self._types            = {}
0030         self._maxLength        = 0
0031         self._tags             = {}
0032         self._tagOrder         = []
0033         self._noCommaSplit     = {}
0034         self._noDefaultClear   = {}
0035         self._setDuringParsing = {}
0036         self._currentlyParsing = False
0037         # now play with the rest
0038         for arg in args:
0039             lower = arg.lower()
0040             if lower == 'python':
0041                 self.register ('storePrepend',
0042                                '',
0043                                VarParsing.multiplicity.singleton,
0044                                VarParsing.varType.string,
0045                                "Prepend location of files starting "
0046                                "with '/store'")
0047                 # Don't continue because we want to process the next
0048                 # piece, too.
0049             if lower == 'analysis' or lower == 'python':
0050                 # Optionos for cmsRun or FWLite.Python
0051                 self.register ('maxEvents',
0052                                -1,
0053                                VarParsing.multiplicity.singleton,
0054                                VarParsing.varType.int,
0055                                "Number of events to process (-1 for all)")
0056                 self.register ('totalSections',
0057                                0,
0058                                VarParsing.multiplicity.singleton,
0059                                VarParsing.varType.int,
0060                                "Total number of sections")
0061                 self.register ('section',
0062                                0,
0063                                VarParsing.multiplicity.singleton,
0064                                VarParsing.varType.int,
0065                                "This section (from 1..totalSections inclusive)")
0066                 self.register ('inputFiles',
0067                                '',
0068                                VarParsing.multiplicity.list,
0069                                VarParsing.varType.string,
0070                                "Files to process")
0071                 self.register ('secondaryInputFiles',
0072                                '',
0073                                VarParsing.multiplicity.list,
0074                                VarParsing.varType.string,
0075                                "Second group of files to process (if needed)")
0076                 self.register ('filePrepend',
0077                                '',
0078                                VarParsing.multiplicity.singleton,
0079                                VarParsing.varType.string,
0080                                "String to prepend location of all files")
0081                 self.register ('outputFile',
0082                                'output.root',
0083                                VarParsing.multiplicity.singleton,
0084                                VarParsing.varType.tagString,
0085                                "Name of output file (if needed)")
0086                 self.register ('secondaryOutputFile',
0087                                '',
0088                                VarParsing.multiplicity.singleton,
0089                                VarParsing.varType.tagString,
0090                                "Name of second output file (if needed)")
0091                 self.register ('tag',
0092                                '',
0093                                VarParsing.multiplicity.singleton,
0094                                VarParsing.varType.string,
0095                                "tag to add to output filename")
0096                 self.setupTags (tag = 'numEvent%d',
0097                                 ifCond = 'maxEvents > 0',
0098                                 tagArg = 'maxEvents')
0099                 self.setupTags (tag = '%s',
0100                                 ifCond = 'tag',
0101                                 tagArg = 'tag')
0102                 continue
0103             # old, depricated, but here for compatibility of older code
0104             if lower == "standard":
0105                 # load in old standard arguments and defaults
0106                 self.register ('maxEvents',
0107                                -1,
0108                                VarParsing.multiplicity.singleton,
0109                                VarParsing.varType.int,
0110                                "Number of events to process (-1 for all)")
0111                 self.register ('files',
0112                                '',
0113                                VarParsing.multiplicity.list,
0114                                VarParsing.varType.string,
0115                                "Files to process")
0116                 self.register ('secondaryFiles',
0117                                '',
0118                                VarParsing.multiplicity.list,
0119                                VarParsing.varType.string,
0120                                "Second group of files to process (if needed)")
0121                 self.register ('output',
0122                                'output.root',
0123                                VarParsing.multiplicity.singleton,
0124                                VarParsing.varType.tagString,
0125                                "Name of output file (if needed)")
0126                 self.register ('secondaryOutput',
0127                                '',
0128                                VarParsing.multiplicity.singleton,
0129                                VarParsing.varType.tagString,
0130                                "Name of second output file (if needed)")
0131                 self.setupTags (tag = 'numEvent%d',
0132                                 ifCond = 'maxEvents > 0',
0133                                 tagArg = 'maxEvents')
0134                 continue
0135             # if we're still here, then we've got a rogue arument
0136             print("Error: VarParsing.__init__ doesn't understand '%s'" \
0137                   % arg)
0138             raise RuntimeError("Failed to create VarParsing object")
0139
0140
0141     def setupTags (self, **kwargs):
0142         """Sets up information for tags for output names"""
0143         necessaryKeys = set (['ifCond', 'tag'])
0144         allowedKeys   = set (['tagArg'])
0145         for key in kwargs.keys():
0146             if key in allowedKeys:
0147                 continue
0148             if key in necessaryKeys:
0149                 necessaryKeys.remove (key)
0150                 continue
0151             # if we're here, then we have a key that's not understood
0152             print("Unknown option '%s'" % key)
0153             raise RuntimeError("Unknown option")
0154         if necessaryKeys:
0155             # if this is not empty, then we didn't have a key that was
0156             # necessary.
0157             print("Missing keys: %s" % necessaryKeys)
0158             raise runtimeError("Missing keys")
0159         tag = kwargs.get('tag')
0160         del kwargs['tag']
0161         self._tags[tag] = kwargs
0162         self._tagOrder.append (tag)
0163
0164
0165     def parseArguments (self):
0166         """Parses command line arguments.  Parsing starts just after
0167         the name of the configuration script.  Parsing will fail if
0168         there is not 'xxxx.py'"""
0169         self._currentlyParsing = True
0170         foundPy      = False
0171         printStatus  = False
0172         help         = False
0173         singleAssign = True
0174         for arg in sys.argv:
0175             if not foundPy and arg.endswith ('.py'):
0176                 foundPy = True
0177                 continue
0178             if not foundPy:
0179                 continue
0180             # If we're here, then we can parse to our hearts content.
0181             # So, is this a command or a declaration?
0182             if arg.count('='):
0183                 # declaration
0184                 name, value = arg.split ('=', 1)
0185                 if name.count('_'):
0186                     # name with command
0187                     name, command = name.split ('_', 1)
0188                     command = command.lower()
0189                     if command == 'load':
0190                         self.loadFromFile (name, value)
0191                         continue
0192                     if command == 'clear':
0193                         self.clearList (name)
0194                         continue
0195                     # If we're here, then I don't recognize this command
0196                     print("Unknown command '%s' in '%s_%s" % \
0197                           (command, name, command))
0198                     raise RuntimeError("Illegal parsing command")
0199                 else:
0200                     # just a name and value
0201                     if name not in self._register:
0202                         print("Error:  '%s' not registered." \
0203                               % name)
0204                         raise RuntimeError("Unknown variable")
0205                     if VarParsing.multiplicity.singleton == \
0206                            self._register[name]:
0207                         # singleton
0208                         if self._beenSet.get (name) and singleAssign:
0209                             print("Variable '%s' assigned multiple times. Use" \
0210                                   , "'multipleAssign' command to avoid")
0211                             raise RuntimeError("Multiple assignment")
0212                         self._beenSet[name] = True
0213                         self.setDefault (name, value)
0214                     else:
0215                         # list
0216                         self.setDefault (name, value)
0217             else:
0218                 # commands
0219                 if arg.count('_'):
0220                     # name modifier
0221                     name, command = arg.split ('_', 1)
0222                     command = command.lower()
0223                     if name not in self._register:
0224                         print("Error:  '%s' not registered." \
0225                               % name)
0226                         raise RuntimeError("Unknown variable")
0227                     if command == 'clear':
0228                         self.clearList (name)
0229                         continue
0230                     # if we're still here, complain that we don't
0231                     # understand this command:
0232                     print("Do not understand '%s' in '%s'" % (command, arg))
0233                     raise RuntimeError("Unknown command")
0234                 else:
0235                     # simple command
0236                     command = arg.lower()
0237                     if command == 'help' or command == '--help':
0238                         help = True
0239                     elif command == 'print' or command == '--print':
0240                         printStatus = True
0241                     elif command == 'noprint' or command == '--noprint':
0242                         printStatus = False
0243                     else:
0244                         # We don't understand this command
0245                         print("Do not understand command '%s'" % (arg))
0246                         raise RuntimeError("Unknown command")
0247             # else if declaration
0248         ###########################
0249         # Post-loading processing #
0250         ###########################
0251         # sections
0252         if 'totalSections' in self._register and \
0253            'section' in self._register and \
0254            'inputFiles' in self._register and \
0255            self.totalSections and self.section:
0256             # copy list
0257             oldInputFiles = self.inputFiles
0258             # clear list
0259             self.clearList ('inputFiles')
0260             # used old list to make list
0261             self.inputFiles = sectionNofTotal (oldInputFiles,
0262                                                self.section,
0263                                                self.totalSections)
0264         # storePrepend
0265         if 'storePrepend' in self._register and \
0266            'inputFiles' in self._register and \
0267            self.storePrepend:
0268             storeRE = re.compile (r'^/store/')
0269             newFileList = []
0270             for filename in self.inputFiles:
0271                 if storeRE.match (filename):
0272                     filename = self.storePrepend + filename
0273                 newFileList.append (filename)
0274             # clear old list
0275             self.clearList ('inputFiles')
0276             # set new list as list
0277             self.inputFiles = newFileList
0278         # filePrepend
0279         if 'filePrepend' in self._register and \
0280            'inputFiles' in self._register and \
0281            self.filePrepend:
0282             newFileList = []
0283             for filename in self.inputFiles:
0284                 filename = self.filePrepend + filename
0285                 newFileList.append (filename)
0286             # clear old list
0287             self.clearList ('inputFiles')
0288             # set new list as list
0289             self.inputFiles = newFileList
0290         # make sure found the py file
0291         if not foundPy:
0292             print("VarParsing.parseArguments() Failure: No configuration " + \
0293                   "file found ending in .py.")
0294             raise RuntimeError("Invalid configuration ending")
0295         if help:
0296             self.help()
0297         if printStatus:
0298             print(self)
0299         self._currentlyParsing = False
0300
0301
0302     def clearList (self, name):
0303         """Empties all entries from list"""
0304         if name not in self._register:
0305             print("Error:  '%s' not registered." \
0306                   % name)
0307             raise RuntimeError("Unknown variable")
0308         if self._register[name] != VarParsing.multiplicity.list:
0309             print("Error: '%s' is not a list" % name)
0310             raise RuntimeError("Faulty 'clear' command")
0311         # if we're still here, do what we came to do
0312         self._lists[name] = []
0313
0314
0315     def setNoDefaultClear (self, name, value=True):
0316         """Tells lists to not clear default list values when set from
0317         command line."""
0318         if name not in self._register:
0319             print("Error:  '%s' not registered." \
0320                   % name)
0321             raise RuntimeError("Unknown variable")
0322         if self._register[name] != VarParsing.multiplicity.list:
0323             print("Error: '%s' is not a list" % name)
0324             raise RuntimeError("Faulty 'setNoDefaultClear' command")
0325         self._noDefaultClear[name] = bool (value)
0326
0327
0328     def setNoCommaSplit (self, name, value=True):
0329         """Tells lists to not split up values by commas."""
0330         if name not in self._register:
0331             print("Error:  '%s' not registered." \
0332                   % name)
0333             raise RuntimeError("Unknown variable")
0334         if self._register[name] != VarParsing.multiplicity.list:
0335             print("Error: '%s' is not a list" % name)
0336             raise RuntimeError("Faulty 'setNoCommaSplit' command")
0337         self._noCommaSplit[name] = bool (value)
0338
0339
0340     def loadFromFile (self, name, filename):
0341         """Loads a list from file"""
0342         if name not in self._register:
0343             print("Error:  '%s' not registered." \
0344                   % name)
0345             raise RuntimeError("Unknown variable")
0346         if self._register[name] != VarParsing.multiplicity.list:
0347             print("Error: '%s' is not a list" % name)
0348             raise RuntimeError("'load' only works for lists")
0349         filename = os.path.expanduser (filename)
0350         if not os.path.exists (filename):
0351             print("Error: '%s' file does not exist.")
0352             raise RuntimeError("Bad filename")
0353         source = open (filename, 'r')
0354         for line in source.readlines():
0355             line = re.sub (r'#.+$', '', line) # remove comment characters
0356             line = line.strip()
0357             if len (line):
0358                 self._lists[name].append( self._convert (name, line ) )
0359         source.close()
0360
0361
0362     def help (self):
0363         """Prints out help information and exits"""
0364         print(self)
0365         print("""Options:
0366         help           : This screen
0367         multipleAssign : Allows singletons to have multiple assignments
0368         print          : Prints out current values
0369         XXX_clear      : Clears list named 'XXX'
0370         """)
0371         sys.exit (0)
0372
0373
0374     def register (self, name,
0375                   default = "",
0376                   mult    = multiplicity.singleton,
0377                   mytype  = varType.int,
0378                   info    = "",
0379                   **kwargs):
0380         """Register a variable"""
0381         # is type ok?
0382         if not VarParsing.multiplicity.isValidValue (mult):
0383             print("Error: VarParsing.register() must use ",\
0384                   "VarParsing.multiplicity.")
0385             raise RuntimeError("Improper 'mult' value")
0386         if not VarParsing.varType.isValidValue (mytype):
0387             print("Error: VarParsing.register() must use ",\
0388                   "VarParsing.varType.")
0389             raise RuntimeError("Improper 'type' value %s" % mytype)
0390         if VarParsing.multiplicity.list == mult and \
0391            VarParsing.varType.tagString == mytype:
0392             print("Error: 'tagString' can only be used with 'singleton'")
0393             raise RuntimeError("Improper registration")
0394         # is the name ok
0395         if name.count ("_"):
0396             print("Error: Name can not contain '_': %s" % name)
0397             raise RuntimeError("Improper 'name'")
0398         # has this been registered before?
0399         if name in self._register:
0400             # Uh oh
0401             print("Error: You can not register a name twice, '%s'" \
0402                   % name)
0403             raise RuntimeError("Attempt to re-register variable")
0404         self._register[name] = mult
0405         self._beenSet[name]  = False
0406         self._info[name]     = info
0407         self._types[name]    = mytype
0408         if len (name) > self._maxLength:
0409             self._maxLength = len (name)
0410         if VarParsing.multiplicity.singleton == mult:
0411             self._singletons[name] = default
0412         else:
0413             self._lists[name] = []
0414             # if it's a list, we only want to use the default if it
0415             # does exist.
0416             if len (default):
0417                 self._lists[name].append (default)
0418         #######################################
0419         ## Process any additional directives ##
0420         #######################################
0421         # do we want to tell the list to not split command line
0422         # arguments by commas?
0423         if kwargs.get ('noCommaSplit'):
0424             self._noCommaSplit[name] = bool( kwargs['noCommaSplit'] )
0425             del kwargs['noCommaSplit']
0426         if kwargs.get ('noDefaultClear'):
0427             self._noDefaultClear[name] = bool( kwargs['noDefaultClear'] )
0428             del kwargs['noDefaultClear']
0429         if len (kwargs):
0430             raise RuntimeError("register() Unknown arguments %s" % kwargs)
0431
0432
0433     def has_key (self, key):
0434         """Returns true if a key is registered"""
0435         return key in self._register
0436
0437
0438     def setType (self, name, mytype):
0439         """Change the type of 'name' to 'mytype'"""
0440         if not VarParsing.varType.isValidValue (mytype):
0441             print("Error: VarParsing.setType() must use ",\
0442                   "VarParsing.varType.")
0443             raise RuntimeError("Improper 'type' value")
0444         oldVal = self.__getattr__ (name, noTags = True)
0445         self._types[name] = mytype
0446         self.setDefault (name, oldVal)
0447
0448
0449     def setDefault (self, name, *args):
0450         """Used to set or change the default of an already registered
0451         name"""
0452         # has this been registered?
0453         if name not in self._register:
0454             print("Error: VarParsing.setDefault '%s' not already registered." \
0455                   % name)
0456             raise RuntimeError("setDefault without registration")
0457         if VarParsing.multiplicity.singleton == self._register[name]:
0458             # make sure we only have one value
0459             if len (args) != 1:
0460                 print("Error: VarParsing.setDefault needs exactly 1 ",\
0461                       "value for '%s'" % name)
0462                 raise RuntimeError("setDefault args problem")
0463             self._singletons[name] = self._convert (name, args[0])
0464         else:
0465             # If:
0466             #   - We have a list (which we do)
0467             #   - We are now processing command line parsing
0468             #   - It has not yet been set after parsing
0469             #   - We have not explicitly asked for it to be NOT cleared
0470             # Then:
0471             #   - We clear the list
0472             if self._currentlyParsing and \
0473                not self._setDuringParsing.get(name) and \
0474                not self._noDefaultClear.get(name):
0475                 # All four conditions have been satisfied, so let's go
0476                 # ahead and clear this.
0477                 self.clearList (name)
0478             # IF we are currently parsing, then tell people I've set
0479             # this:
0480             if self._currentlyParsing:
0481                 self._setDuringParsing[name] = True
0482             # if args is a tuple and it only has one entry, get rid of
0483             # the first level of tupleness:
0484             if isinstance (args, tuple) and len (args) == 1:
0485                 args = args[0]
0486             # is this still a tuple
0487             if isinstance (args, tuple):
0488                 mylist = list (args)
0489             elif isinstance (args, list):
0490                 mylist = args
0491             else:
0492                 mylist = []
0493                 mylist.append (args)
0494             if not self._noCommaSplit.get (name):
0495                 oldList = mylist
0496                 mylist = []
0497                 for item in oldList:
0498                     mylist.extend( VarParsing.commaRE.split( item ) )
0499             for item in mylist:
0500                 self._lists[name].append( self._convert (name, item ) )
0501
0502
0503     def _convert (self, name, inputVal):
0504         """Converts inputVal to the type required by name"""
0505         inputVal = str (inputVal)
0506         if self._types[name] == VarParsing.varType.bool:
0507             if VarParsing.trueRE.match (inputVal) or '1' == inputVal:
0508                 return True
0509             elif VarParsing.falseRE.match (inputVal) or '0' == inputVal:
0510                 return False
0511             # if we're still here, then we don't have 'true' or 'false'
0512             raise RuntimeError("Unknown bool value '%s'.  Must be 'true' or 'false'" % inputVal)
0513         if self._types[name] == VarParsing.varType.string or \
0514            self._types[name] == VarParsing.varType.tagString:
0515             return inputVal
0516         elif self._types[name] == VarParsing.varType.int:
0517             return int (inputVal, 0)
0518         elif self._types[name] == VarParsing.varType.float:
0519             return float (inputVal)
0520         else:
0521             raise RuntimeError("Unknown varType")
0522
0523
0524     def _withTags (self, name):
0525         if name not in self._register:
0526             print("Error:  '%s' not registered." \
0527                   % name)
0528             raise RuntimeError("Unknown variable")
0529         if self._register[name] == VarParsing.multiplicity.list:
0530             print("Error: '%s' is a list" % name)
0531             raise RuntimeError("withTags() only works on singletons")
0532         retval = self._singletons[name]
0533         if retval.endswith ('.root'):
0534             retval, garbage = os.path.splitext (retval)
0535         reverseOrder = self._tagOrder
0536         reverseOrder.reverse()
0537         for tag in reverseOrder:
0538             tagDict = self._tags[tag]
0539             ifCond = tagDict['ifCond']
0540             if ifCond.count('%'):
0541                 pass
0542             else:
0543                 ifCond = "self." + ifCond
0544             boolValue = eval (ifCond)
0545             tagArg = tagDict.get ('tagArg')
0546             if tagArg:
0547                 evalString = "'%s' %% self.%s" % (tag, tagArg)
0548                 tag = eval (evalString)
0549             if boolValue:
0550                 retval = retval + "_" + tag
0551         return retval + ".root"
0552
0553
0554     def __str__ (self):
0555         """String form of self"""
0556         maxLen = min (self._maxLength, 20)
0557         form     = "  %%-%ds: %%s" % maxLen
0558         formInfo = "  %%%ds  - %%s" % (maxLen - 2)
0559         formItem = "  %%%ds    %%s" % (maxLen - 1)
0560         retval = ""
0561         if len (self._singletons):
0562             retval = retval + "Singletons:\n"
0563         for varName, value in sorted (self._singletons.items()):
0564             retval = retval + form % (varName, value) + "\n";
0565             if self._info.get(varName):
0566                 retval = retval + formInfo % ('', self._info[varName]) + "\n"
0567         if len (self._singletons):
0568             retval = retval +  "Lists:\n"
0569         for varName, value in sorted (self._lists.items()):
0570             stringValue = "%s" % value
0571             if len (stringValue) < 76 - maxLen:
0572                 retval = retval + form % (varName, value) + "\n"
0573             else:
0574                 varLength = len (value)
0575                 for index, item in enumerate (value):
0576                     if index == 0:
0577                         retval = retval + form % (varName, "['" + item)
0578                     else:
0579                         retval = retval + formItem % ('',"'" + item)
0580                     if index == varLength - 1:
0581                         retval = retval + "' ]\n"
0582                     else:
0583                         retval = retval + "',\n"
0584             if self._info.get(varName):
0585                 retval = retval + formInfo % ('', self._info[varName]) + "\n"
0586         return retval
0587
0588
0589     def __setattr__ (self, name, value, *extras):
0590         """Lets me set internal values, or uses setDefault"""
0591         if not name.startswith ("_"):
0592             mylist = list (extras)
0593             mylist.insert (0, value)
0594             self.setDefault (name, *mylist)
0595         else:
0596             object.__setattr__ (self, name, value)
0597
0598
0599     def __getattr__ (self, name, noTags = False):
0600         """Lets user get the info they want with obj.name"""
0601         if name.startswith ("_"):
0602             # internal use
0603             return object.__getattribute__ (self, name)
0604         else:
0605             # user variable
0606             if name not in self._register:
0607                 print("Error:  '%s' not already registered." \
0608                       % name)
0609                 raise RuntimeError("Unknown variable")
0610             if VarParsing.multiplicity.singleton == self._register[name]:
0611                 if VarParsing.varType.tagString == self._types[name] \
0612                        and not noTags:
0613                     return self._withTags (name)
0614                 else:
0615                     return self._singletons[name]
0616             else:
0617                 return self._lists[name]
0618
0619
0620 ##############################################################################
0621 ## ######################################################################## ##
0622 ## ##                                                                    ## ##
0623 ## ######################################################################## ##
0624 ##############################################################################
0625
0626
0627 if __name__ == "__main__":
0628     #############################################
0629     ## Load and save command line history when ##
0630     ## running interactively.                  ##
0631     #############################################
0632     import os, readline
0633     import atexit
0634     historyPath = os.path.expanduser("~/.pyhistory")
0635
0636
0637     def save_history(historyPath=historyPath):
0638         import readline
0639         readline.write_history_file(historyPath)
0640         if os.path.exists(historyPath):
0641             readline.read_history_file(historyPath)
0642
0643
0644     atexit.register(save_history)
0645     readline.parse_and_bind("set show-all-if-ambiguous on")
0646     readline.parse_and_bind("tab: complete")
0647     if os.path.exists (historyPath) :
0648         readline.read_history_file(historyPath)
0649         readline.set_history_length(-1)
0650
0651
0652     ############################
0653     # Example code starts here #
0654     ############################
0655
0656     obj = VarParsing ('standard')
0657     obj.register ('someVar',
0658                   mult=VarParsing.multiplicity.singleton,
0659                   info="for testing")
0660     obj.setupTags (tag    = "someCondition",
0661                    ifCond = "someVar")
0662     ## obj.register ('numbers',
0663     ##               mult=VarParsing.multiplicity.list,
0664     ##               info="Numbers")
0665     ## obj.register ('singleNumber',
0666     ##               mult=VarParsing.multiplicity.singleton,
0667     ##               info="A single number")
0668     obj.parseArguments()
0669