PythonUtilities/scripts/filterCSVwithJSON.py

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