hp_cmd.py 2.26 KB
import sys

from run import generate_csv_file_if_needed, init_console_logger, config, get_default_sta_sto
from datetime import datetime
import argparse

if __name__ == "__main__":
    hp_parser = argparse.ArgumentParser(prog='Heliopropa Command tool')

    hp_parser.add_argument('-clog', '--console-log', action='store_true')
    hp_parser.add_argument('-s', '--dry-run', action='store_true')

    subparsers = hp_parser.add_subparsers(dest='hpcmd_name')

    csvgen_parser = subparsers.add_parser('show_config',
                                          help='Show configuration targets and inputs')
    csvgen_parser = subparsers.add_parser('csv_gen',
                                          help='Generate CSV for given target/input (can be input/input)')
    csvgen_parser.add_argument('-a', '--all', action='store_true')
    csvgen_parser.add_argument('-t', '--target')
    csvgen_parser.add_argument('-i', '--input')
    csvgen_parser.add_argument('-b', '--begin')
    csvgen_parser.add_argument('-e', '--end')

    args = hp_parser.parse_args()

    if args.console_log:
        _logger = init_console_logger()

    if args.hpcmd_name == 'show_config':
        for meteor in config['targets'] + config['inputs']:
            meteor_type = meteor['type']
            meteor_name = f"'{meteor['name']}'"
            meteor_slug = meteor['slug']
            meteor_models = ', '.join(list(meteor['models'].keys()))
            print(f"{meteor_type:6} {meteor_slug:6} {meteor_name:11}: models = {meteor_models}")
    elif args.hpcmd_name == 'csv_gen':
        # build default dates in not -b or not -e
        sta, sto = get_default_sta_sto(args.begin, args.end)

        meteors_list = config['targets'] + config['inputs']

        if args.all:
            # build the {target_slug: [input list]} dictionnary from config.yaml
            meteors_dict = {t['slug']: list(t['models'].keys()) for t in meteors_list}
        else:
            # todo: check -i an -t
            meteors_dict = {args.target: [args.input]}

        for _t, _i_list in meteors_dict.items():
            # print(_t, _i_list)
            for _i in _i_list:
                print(f"generating csv for {_t}/ {_i}")
                if args.dry_run:
                    continue
                generate_csv_file_if_needed(_t, _i, sta, sto)