parset_concat.py

#!/usr/bin/env python3
import sys
import os
import re
import logging
from argparse import ArgumentParser
## TODO add input validation


def parse_arguments(args):
    parser = ArgumentParser(description='Concatenate step defined to the input parset')
    parser.add_argument('step_name', help='name of the step')
    parser.add_argument('step_type', help='type of the step')
    parser.add_argument('--input_parset', default=None, help='input parset')
    parser.add_argument('step_arguments', help='step arguments', nargs='*')

    return parser.parse_args()


def main(args):
    arguments = parse_arguments(args)
    concatenate(arguments.input_parset,
                arguments.step_name,
                arguments.step_type,
                arguments.step_arguments)


def concatenate(lines, step_name, step_type, step_args):
    """
    Append a new step to existing parset. Output is printed to stdout.

    Args:
        lines (List[str]): list of existing lines
        step_name (str): name of new step
        step_type (str): type of new step
        step_args (List[str]): key-value pairs, e.g. ["applybeam=true"]

    Returns:
        None
    """
    steps = []
    if lines:
        for line in lines:
            if line.startswith("steps="):
                _, steps = line.split("=")
                steps = [step for step in steps.strip().lstrip("[").rstrip("]").split(",") if step]
            else:
                print(line)

    print("%s.type=%s" % (step_name, step_type))
    if step_args:
        for arg in step_args:
            if not '=' in arg:
                logging.error('Argument invalid (e.g. key=value): %s', arg)
                sys.exit(1)
            key = '.'.join([step_name, arg.strip()])
            print(key)

    if step_name in steps:
        logging.error('Step named %s already present in steps: %s', step_name, steps)
        exit(1)
    else:
        steps += [step_name]
    print("steps=[%s]" % ",".join(steps))


def concatenate_file(parsetfile, step_name, step_type, args):
    print(args)
    if parsetfile and os.path.exists(parsetfile):
      with open(parsetfile, "r") as f_stream:
          concatenate(f_stream.read().splitlines(), step_name, step_type, args)


if __name__=='__main__':
    main(sys.argv)