playbooks/debian/sa-merge-all.py

   1 #! /usr/bin/env python
   2
   3 import argparse
   4
   5 if __name__ == '__main__':
   6     parser = argparse.ArgumentParser(description='Import')
   7     parser.add_argument('-o', '--output', required=True, help='output file')
   8     parser.add_argument('input_fnames', nargs='+', help='input files to merge')
   9     args = parser.parse_args()
  10
  11     input_files = {}
  12
  13     def close_all():
  14         for f in input_files.values():
  15             f.close()
  16
  17     for fname in args.input_fnames:
  18         try:
  19             input_files[fname] = open(fname, 'rtU')
  20         except (IOError, OSError):
  21             close_all()
  22             raise
  23
  24     # Read and process the first 3 lines
  25     counters = {'spam': {}, 'nonspam': {}}
  26     for fname, infile in input_files.items():
  27         version_line = infile.readline()
  28         for _line_no in 1, 2:
  29             counter_line = infile.readline().strip()
  30             v, counter, tag = counter_line.split('\t')
  31             if v != 'v':
  32                 close_all()
  33                 raise ValueError(
  34                         'Bad v-tag in file %s line %s: unknown v-tag %s, '
  35                         'expected "v", got %r' % (fname, counter_line, v))
  36             if tag == 'num_spam':
  37                 try:
  38                     counters['spam'][fname] = int(counter)
  39                 except ValueError:
  40                     close_all()
  41                     raise
  42             elif tag == 'num_nonspam':
  43                 try:
  44                     counters['nonspam'][fname] = int(counter)
  45                 except ValueError:
  46                     close_all()
  47                     raise
  48             else:
  49                 close_all()
  50                 raise ValueError(
  51                         'Bad tag in file %s line %s: unknown tag %s, '
  52                         'expected "num_spam" or "num_nonspam", got %r' % (
  53                             fname, counter_line, tag))
  54
  55     counters_total = {'spam': 0, 'nonspam': 0}
  56     for key, files in counters.items():
  57         counters_total[key] = sum(files.values())
  58
  59     try:
  60         with open(args.output, 'wt') as outfile:
  61             outfile.write(version_line)
  62             for key, value in counters_total.items():
  63                 outfile.write('v\t%d\tnum_%s\n' % (value, key))
  64             for infile in input_files.values():
  65                 for line in infile:
  66                     outfile.write(line)
  67     except (IOError, OSError):
  68         close_all()
  69         raise
  70     close_all()