metarun.py

#!/usr/bin/env python

"""
Run command varying its arguments
"""

from subprocess import check_call
from multiprocessing import Pool
from itertools import product


def parse_params(paramstrlist, sep):
    """
    Parse list of 'PARAM=p1,p2,p3,...'-type string

    >>> parse_params(['A=1,2,3', 'B=x,y,z'], ',')
    (['A', 'B'], [['1', '2', '3'], ['x', 'y', 'z']])

    """
    val_list = []
    key_list = []
    for paramstr in paramstrlist:
        (key, vals) = paramstr.split('=')
        key_list.append(key)
        val_list.append(vals.split(sep))
    return (key_list, val_list)


def gene_cmds_normal(script, keys, vals, pre='--'):
    """
    Generate list of commands from script name, option names, and sets of value

    >>> cmds = gene_cmds_normal('run', ['A', 'B'], [['1', '2'], ['x', 'y']])
    >>> list(cmds)  #doctest: +NORMALIZE_WHITESPACE
    [['run', '--A 1', '--B 2'],
     ['run', '--A x', '--B y']]
    >>> cmds = gene_cmds_normal('python run.py',
    ...                         ['A', 'B'], [['1', '2'], ['x', 'y']])
    >>> list(cmds)  #doctest: +NORMALIZE_WHITESPACE
    [['python', 'run.py', '--A 1', '--B 2'],
     ['python', 'run.py', '--A x', '--B y']]

    """
    script = script.split()
    for val in vals:
        yield script + ['%s%s %s' % (pre, k, v) for (k, v) in zip(keys, val)]


def gene_cmds_format(script_fmt, keys, vals):
    """
    Generate list of commands from format str, option names, and sets of value

    >>> cmds = gene_cmds_format('run A+B="%(A)s+%(B)s"',
    ...                         ['A', 'B'], [['1', '2'], ['x', 'y']])
    >>> list(cmds)  #doctest: +NORMALIZE_WHITESPACE
    [['run', 'A+B="1+2"'],
     ['run', 'A+B="x+y"']]
    >>> cmds = gene_cmds_format('python run.py A+B="%(A)s+%(B)s"',
    ...                         ['A', 'B'], [['1', '2'], ['x', 'y']])
    >>> list(cmds)  #doctest: +NORMALIZE_WHITESPACE
    [['python', 'run.py', 'A+B="1+2"'],
     ['python', 'run.py', 'A+B="x+y"']]

    """
    for val in vals:
        yield (script_fmt % dict(zip(keys, val))).split()


def run(cmd):
    print ' '.join(cmd)
    check_call(cmd)


def dry_run(cmd):
    print ' '.join(cmd)


def main():
    from optparse import OptionParser
    parser = OptionParser(
        usage='%prog [options] SCRIPT PARAM_A=a1,a2,a3 [PARAM_B=b1,b2 ...]',
        description=__doc__)
    parser.add_option(
        "-r", "--pre", default='--',
        help="prefix for your command line options (default: '--')")
    parser.add_option(
        "-s", "--sep", default=',',
        help="separator for the parameter values (default: ',')")
    parser.add_option(
        "-f", "--format", default=False, action="store_true",
        help=("use python string format such as '%(PARAM_A)s' for "
              "the SCRIPT argument"))
    parser.add_option(
        "-n", "--dry-run", default=False, action="store_true",
        help="don't actually do anything, just print commands")
    parser.add_option(
        "-p", "--processes", default=1, type=int,
        help="number of process to use (default: 1)")
    (opts, args) = parser.parse_args()

    script = args[0]
    (keys, val_list) = parse_params(args[1:], opts.sep)
    vals = product(*val_list)

    if opts.format:
        cmds = gene_cmds_format(script, keys, vals)
    else:
        cmds = gene_cmds_normal(script, keys, vals, opts.pre)

    if opts.processes > 1:
        pool = Pool(processes=opts.processes)
        _map = pool.map
    else:
        _map = map

    if opts.dry_run:
        _map(dry_run, cmds)
    else:
        _map(run, cmds)


if __name__ == '__main__':
    main()