testcli_argparse.py 2.82 KB
Newer Older
Vlastimil Zima's avatar
Vlastimil Zima committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72
#!/usr/bin/python3
"""Test client interface."""
import argparse
import json
from datetime import date, datetime
from typing import Sequence, cast

import testcli


def parse_date(value: str):
    return datetime.strptime(value, '%Y-%m-%d').date()


def parse_params(value: str):
    return value.split('=', 1)


def get_parser():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=['create', 'drop'])
    parser.add_argument('label')
    parser.add_argument('--version', action='version', version=testcli.__version__)
    parser.add_argument('--config', type=argparse.FileType('r'), help="Set config file.")
    parser.add_argument('--string', help="Set custom string.")
    parser.add_argument('--number', type=int, help="Set custom number.")
    parser.add_argument('--params', type=parse_params, action='append', help="Set custom parameters.")
    parser.add_argument('--today', type=parse_date, default=date.today(), help="Set custom today [default: today].")
    # Can't use `default=['all']`, because arguments are appended to the default.
    parser.add_argument('--choices', action='append', type=testcli.Choice, choices=[c.value for c in testcli.Choice],
                        help="Set custom choices [default: all]. Available options: %(choices)s.")
    parser.add_argument('-n', '--dry-run', action='store_true', default=False, help="Don't actually do anything.")
    parser.add_argument('-v', '--verbosity', type=int, default=1,
                        help="Set verbosity level in range 0 to 3 [default: %(default)s].")
    parser.set_defaults(**{k: v for k, v in testcli.DEFAULT_CONFIG.items() if k != 'params'})
    return parser


def main(argv: Sequence[str] = None):
    parser = get_parser()
    options = parser.parse_args(argv)

    if options.config:
        config = json.loads(options.config.read())
        parser.set_defaults(**{k: v for k, v in config.items() if k != 'params'})

    # Parse again to consider config.
    options = parser.parse_args(argv)

    if not options.choices or testcli.Choice.ALL in options.choices:
        choices = [c for c in testcli.Choice if c != testcli.Choice.ALL]
    else:
        choices = options.choices

    if options.params:
        params = dict(options.params)
    elif options.config:
        params = config.get('params', testcli.DEFAULT_CONFIG['params'])
    else:
        params = testcli.DEFAULT_CONFIG['params']

    if options.command == 'create':
        testcli.create(options.label, options.string, options.number, params, options.today, choices, options.dry_run,
                       options.verbosity)
    else:
        assert options.command == 'drop'
        testcli.drop(options.label, options.string, options.number, params, options.today, choices, options.dry_run,
                     options.verbosity)


if __name__ == '__main__':
    main()