pstat

#!/usr/bin/env python

import collections

import qstatpretty.config
import qstatpretty.source
import qstatpretty.pretty
import qstatpretty.parser
import qstatpretty.ttyutil.table as ttytable
import qstatpretty.ttyutil.resize as ttyresize
import qstatpretty.passthrough_optparse as ptoptparse


EPILOG = """
pstat is a wrapper for the qstat command of Grid Engine / Torque Resource Manager.
It takes the XML output from qstat and pretty-prints it.

All additional non-pstat options will be passed to qstat:

    $ pstat -u johndoe

"""


def parse_args():
    def cb_source(option, opt, value, parser, *args):
        parser.values.source = args[0]

        if len(args) == 2:
            setattr(parser.values, args[1], value)

    parser = ptoptparse.PassThroughOptionParser(version="0.1.0", epilog=EPILOG)
    parser.add_option(
        "--flavor",
        dest="flavor",
        choices=list(qstatpretty.parser.get_parser_names()),
        help="Choose a flavor of grid system (choices: {choices}) [default: \"%default\"]".format(
            choices=", ".join(qstatpretty.parser.get_parser_names())
        )
    )

    parser.add_option(
        "--delimiters",
        dest="delimiters",
        choices=list(ttytable.DELIMITERS.keys()),
        help="Choose a delimiter style for displaying the formatted table (choices: {choices}) [default: \"%default\"]".format(
            choices=", ".join(ttytable.DELIMITERS.keys())
        )
    )

    parser.add_option(
        "--table",
        dest="table_algorithm",
        choices=list(ttyresize.TABLE_ALGORITHMS.keys()),
        help="Choose an algorithm for fitting the table to the terminal (choices: {choices}) [default: \"%default\"]".format(
            choices=", ".join(ttyresize.TABLE_ALGORITHMS.keys())
        )
    )

    parser.add_option(
        "--source-local",
        dest="source",
        action="callback",
        callback=cb_source,
        callback_args=("local",),
        nargs=0,
        help="Get job status from a locally running grid manager"
    )

    parser.add_option(
        "--source-ssh",
        dest="source",
        action="callback",
        callback=cb_source,
        callback_args=("ssh", "source_ssh_hostname"),
        nargs=1,
        type=str,
        metavar="HOSTNAME",
        help="Get job status by querying via SSH to HOSTNAME"
    )

    parser.add_option(
        "--source-file",
        dest="source",
        action="callback",
        callback=cb_source,
        callback_args=("file", "source_file_path"),
        nargs=1,
        type=str,
        metavar="FILENAME",
        help="Get job status by parsing FILENAME"
    )

    parser.set_defaults(**qstatpretty.config.get_config())
    opt, args = parser.parse_args()

    if not opt.table_algorithm:
        parser.error("No --table specification!")

    if not opt.delimiters:
        parser.error("No --delimiters specification!")

    return opt, args


def main():
    opt, args = parse_args()

    source = qstatpretty.source.SOURCES[opt.source](opt)

    parser = qstatpretty.parser.get_parser_by_name(opt.flavor)
    jobs = parser.parse(source(parser, args))

    job_states = collections.defaultdict(int)
    for job in jobs:
        job_states[job['state']] += 1

    qstatpretty.pretty.pretty_table(jobs, parser.table_columns, table_algorithm=ttyresize.TABLE_ALGORITHMS[opt.table_algorithm], delimiters=ttytable.DELIMITERS[opt.delimiters])

    # Setup summary message
    summary_msg = "{n_jobs} job{plural} total"
    summary = {
        'n_jobs': len(jobs),
        'plural': 's' if len(jobs) != 1 else ''
    }

    if len(jobs) > 0:
        summary_msg = summary_msg + " ({states})."
        summary['states'] = ", ".join(
            "{0}: {1}".format(k, job_states[k])
            for k in sorted(job_states.keys())
        )

    # Print qstat summary
    print(summary_msg.format(**summary))


if __name__ == '__main__':
    main()