update_config.py

import argparse
import json
import sys

import yaml
import zmq

# Defining available options
agipd_options = {
        "force-hg-if-below": {'typ': int},
        "rel-gain": {'typ': bool},
        "xray-gain": {'typ': bool},
        "blc-noise": {'typ': bool},
        "blc-set-min": {'typ': bool},
        "dont-zero-nans": {'typ': bool},
        "dont-zero-orange": {'typ': bool},
        "max-pulses": {'typ': list,
                       'msg': "Range list of maximum pulse indices "
                              "(--max-pulses start end step). "
                              "3 max input elements. "},
        "calfile": {'typ': str},
    }

data_mapping = {
        "karabo-da": {'typ': list,
                      'choices': ["AGIPD{:02d}".format(i) for i in range(16)],
                      'msg': "Choices: [AGIPD00 ... AGIPD15]. "}
}

available_options = {
    "SPB_DET_AGIPD1M-1": [agipd_options, data_mapping],
    "MID_DET_AGIPD1M-1": [agipd_options, data_mapping]
}

formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52)
parser = argparse.ArgumentParser(
    description='Request update of configuration', formatter_class=formatter)
required_args = parser.add_argument_group('required arguments')
required_args.add_argument('--karabo-id', type=str,
                           choices=['SPB_DET_AGIPD1M-1',
                                    'MID_DET_AGIPD1M-1'])
required_args.add_argument('--proposal', type=str,
                           help='The proposal number, without leading p, '
                                'but with leading zeros.')
required_args.add_argument('--cycle', type=str, help='The facility cycle.')
parser.add_argument('--apply', action='store_true',
                    help='Apply and push the requested '
                         'configuration update to the git.')
# remove help calls as they will cause the argument parser to exit
add_help = False
if "-h" in sys.argv:
    sys.argv.remove("-h")
    add_help = True
if "--help" in sys.argv:
    sys.argv.remove("--help")
    add_help = True

known, remaining = parser.parse_known_args()
args = vars(known)
karabo_id = args["karabo_id"]
bool_keys = []


# Avoid erros when karabo_id not yet given through the command line.
if karabo_id is not None:
    # adding "no" bools to available options
    for det, val in available_options[karabo_id][0].items():
        if val['typ'] == bool:
            bool_keys.append(det)

    for b in bool_keys:
        available_options[karabo_id][0]['no-{}'.format(det)] = {'typ': bool}

    for exposed_options in available_options[karabo_id]:
        for option, info in exposed_options.items():
            metavar = None
            if info['typ'] == list:
                nargs = '+'
                typ = str
                action = 'append'
                # Avoid having a big line of choices in the help message.
                if 'choices' in info.keys():
                    metavar = option.upper()
            else:
                action = None
                nargs = None
                typ = info['typ']
            # Add help messages
            help_msg = ""
            if 'msg' in info.keys():
                help_msg += info['msg']
            if 'choices' in info.keys():
                choices = info['choices']
            else:
                choices = None
            help_msg += f"Type: {info['typ'].__name__} ".upper()
            parser.add_argument(f"--{option}", type=typ, action=action,
                                metavar=metavar, nargs=nargs, choices=choices,
                                help=help_msg)

parser.add_argument('--instrument', type=str, choices=["CALLAB"],
                    help='This is only used for testing purposes.')

if add_help:
    sys.argv.append("--help")

args = vars(parser.parse_args())

task = "correct"
# check if instrument is not given from CLI (e.g. CALLAB)
if args['instrument'] is None:
    # extract instrument from karabo_id
    instrument = karabo_id.split("_")[0]
else:
    instrument = args['instrument']

proposal = args['proposal']
cycle = args['cycle']

if instrument is None or proposal is None or cycle is None:
    print("Need to define all fields")
    exit()

new_conf = {task: {instrument: {karabo_id: {}}}}
for key, value in args.items():
    key = key.replace("_", "-")
    for exposed_options in available_options[karabo_id]:
        if key in exposed_options and value is not None:

            if isinstance(value, list):
                value = value[0]
            # convert no arguments to bool false
            if 'no-' in key and isinstance(value, bool):
                if key not in bool_keys:
                    new_conf[task][instrument][karabo_id][key.replace('no-', '')] = False  # noqa
                # avoid saving the "no-"key into the updated config
                continue
            # Assure adding an empty string for new empty
            # str. updates (e.g. calfile)
            if isinstance(key, str) and (value == '' or value == ' '):
                value = '""'
            # checking if data-mapping was updated.
            if key in data_mapping.keys():
                if 'data-mapping' not in new_conf.keys():
                    new_conf['data-mapping'] = {karabo_id: {key: {}}}
                new_conf['data-mapping'][karabo_id][key] = value
            else:
                new_conf[task][instrument][karabo_id][key] = value

pyaml = yaml.dump(new_conf, default_flow_style=False)

if not args["apply"]:
    print("\n")
    print("-" * 80)
    print("THIS IS A DRY RUN ONLY, NO CHANGES ARE MADE")
    print("\n")
    print("-" * 80)

print(f"Sending the following update: \n {pyaml}")
print("-" * 80)
con = zmq.Context()
socket = con.socket(zmq.REQ)
con = socket.connect("tcp://max-exfl016:5555")
msg = "','".join(["update_conf", "SASEX", args["karabo_id"],
                  instrument, args["cycle"], args["proposal"],
                  json.dumps(new_conf), str(args["apply"])])
socket.send("['{}']".format(msg).encode())
resp = socket.recv_multipart()[0]
print("Configuration now in place is:")
print(resp.decode())