PythonUtilities/scripts/filterCSVwithJSON.py

0001 #!/usr/bin/env python3
0002
0003 from __future__ import print_function
0004 import sys
0005 import optparse
0006 import re
0007 from FWCore.PythonUtilities.LumiList import LumiList
0008
0009
0010 if __name__ == '__main__':
0011
0012     parser = optparse.OptionParser ("Usage: %prog input.json input.csv output.csv")
0013     parser.add_option ('--output', dest='output', type='string',
0014                        help='Save output to file OUTPUT')
0015     parser.add_option ('--runIndex', dest='runIndex', type='int',
0016                        default = 0,
0017                        help='column to be converted to run number (default %default)')
0018     parser.add_option ('--lumiIndex', dest='lumiIndex', type='int',
0019                        default = 1,
0020                        help='column to be converted to lumi section number (default %default)')
0021     parser.add_option ('--noWarnings', dest='noWarnings', action='store_true',
0022                        help='do not print warnings about lines not matching run, lumi numbers')
0023     # required parameters
0024     (options, args) = parser.parse_args()
0025     if len (args) != 3:
0026         raise RuntimeError("Must provide an input JSON file, an input CSV file, and an output CSV file")
0027
0028     sepRE = re.compile (r'[\s,;:]+')
0029     runLumiDict = {}
0030     jsonList = LumiList (args[0])
0031     source = open (args[1], 'r')
0032     target = open (args[2], 'w')
0033     runIndex, lumiIndex = options.runIndex, options.lumiIndex
0034     minPieces = max (runIndex, lumiIndex) + 1
0035     for line in source:
0036         copy = line.strip()
0037         pieces = sepRE.split (copy.strip())
0038         if len (pieces) < minPieces:
0039             if not options.noWarnings:
0040                 print("Saving line '%s' since no identifiable run and lumi info" \
0041                       % copy)
0042             target.write (line)
0043             continue
0044         try:
0045             run, lumi = int( pieces[runIndex] ), int( pieces[lumiIndex] )
0046         except:
0047             if not options.noWarnings:
0048                 print("Saving line '%s' since no identifiable run,lumi info" \
0049                       % copy)
0050             target.write (line)
0051             continue
0052         # OK.  We recognize this line as containing a valid run and
0053         # lumi number.  Is it part of the JSON file we provided?
0054         if (run, lumi) in jsonList:
0055             # Yes, it is.  Take it!
0056             target.write (line)
0057
0058     source.close()
0059     target.close()