LHEInterface/scripts/addLHEnumbers.py

0001 #!/usr/bin/env python3
0002
0003 from __future__ import print_function
0004 import logging
0005 import argparse
0006 import sys
0007 import os
0008 import re
0009
0010
0011 def number_events(input_file, output_file=None, offset=0):
0012     if output_file is None:
0013         output_file = input_file
0014     if not os.path.exists(os.path.dirname(os.path.realpath(output_file))):
0015         os.makedirs(os.path.dirname(os.path.realpath(output_file)))
0016
0017     nevent = offset
0018     with open('tmp.txt', 'w') as fw:
0019         with open(input_file, 'r') as ftmp:
0020             for line in ftmp:
0021                 if re.search('\s*</event>', line):
0022                     nevent += 1
0023                     fw.write('<event_num num="' + str(nevent) +  '"> ' + str(nevent) + '</event_num>\n')
0024                 fw.write(line)
0025     if output_file is not None:
0026         os.rename("tmp.txt", output_file)
0027     else:
0028         os.rename("tmp.txt", input_file)
0029     return nevent
0030
0031
0032 if __name__=="__main__":
0033
0034     parser = argparse.ArgumentParser(
0035         description="Add numbers to lhe")
0036     parser.add_argument("input_file", type=str,
0037                         help="Input LHE file path.")
0038     parser.add_argument("-o", "--output-file", default=None, type=str,
0039                         help="Output LHE file path. If not specified, output to input file")
0040     args = parser.parse_args()
0041
0042     logging.info('>>> launch addLHEnumbers.py in %s' % os.path.abspath(os.getcwd()))
0043
0044     logging.info('>>> Input file: [%s]' % args.input_file)
0045     logging.info('>>> Write to output: %s ' % args.output_file)
0046
0047     number_events(args.input_file, args.output_file)