chop

#!/usr/bin/env python
import argparse
from datetime import datetime
from functools import partial
import os
import re

from progressbar import ProgressBar, Percentage, Bar, ETA

"""
Extract entries in a time interval from a log file.
"""


def parse_commandline_args():
    parser = argparse.ArgumentParser(
        description='Extract entries in a time interval from a log file.')
    parser.add_argument('--start', action='store', dest='start_time',
                        default="2000-01-01 01:00:00.000",
                        help='Log start timestamp')
    parser.add_argument('--end', action='store', dest='end_time',
                        default="3000-12-31 01:00:00.000",
                        help='Log end timestamp')
    parser.add_argument('--dest', action='store', dest='dest',
                        default=None,
                        help=('Destination directory. '
                              'Defaults to directories of each logfile, resp.'
                              ))
    parser.add_argument('--regex', action='store', dest='regex',
                        default=None,
                        help=('Regex pattern for matching ' +
                              'first line in each log entry'))
    parser.add_argument('-q', '--quiet', action='store_true', dest='quiet',
                        default=False,
                        help='Don\'t show progressbar')
    parser.add_argument('files', metavar='FILE', type=str, nargs='+',
                        help='Log file(s) to be chopped.')
    return parser.parse_args()


def count_lines(filename):
    return sum(1 for line in open(filename, 'r'))


def process_line(line, timestamp_processor):
    try:
        return True, timestamp_processor(line[:23])
    except:
        return False, None


def main():
    # parse command line args
    args = parse_commandline_args()

    # setup timestamp converter
    LOG_PATTERN = "%Y-%m-%d %H:%M:%S.%f"
    str_to_datetime = partial(
        lambda x, y: datetime.strptime(y, x),
        LOG_PATTERN,
        )

    # convert timestamp bookends to datetime objects
    start_time = str_to_datetime(args.start_time)
    end_time = str_to_datetime(args.end_time)

    # setup RegExp matcher for matching first line of each entry
    if args.regex is not None:
        _regex = re.compile(args.regex)
    else:
        _regex = None

    for filename in args.files:
        print '***', filename, '***'

        num_lines = count_lines(filename)

        dest = args.dest or os.path.dirname(os.path.abspath(filename))
        ext = '' if args.dest is not None else '.chopped'
        outfile = os.path.join(
                    dest,
                    os.path.basename(filename) + ext)

        with open(filename, 'r') as f:
            with open(outfile, 'w') as o:
                entry = ''
                if not args.quiet:
                    bar = ProgressBar(
                        widgets=[Percentage(), Bar(), ETA()],
                        maxval=num_lines*1.2,
                        ).start()
                for ix, line in enumerate(f):
                    starts_log, dt = process_line(
                        line, str_to_datetime)

                    if starts_log and \
                       dt >= start_time and \
                       dt <= end_time:
                        # skip the log entry if the first line doesn't match
                        # the specified regex pattern
                        if _regex is not None and \
                           not _regex.search(line):
                            continue

                        # write the previous log entry to the file
                        if entry != '':
                            o.write(entry)
                            entry = ''

                        # add the first line
                        entry += line

                    elif not starts_log and entry != '':
                        # add subsequent line (multi-line log entry)
                        entry += line

                    if not args.quiet:
                        bar.update(ix)

                if entry != '':
                    o.write(entry)

                if not args.quiet:
                    bar.finish()
        print


if __name__ == '__main__':
    main()