path: root/config.py

#!/usr/bin/env python3

"""Global configuration driven by commandline arguments (even across
different modules).  Usage:

    module.py:
    ----------
    import config

    parser = config.add_commandline_args(
        "Module",
        "Args related to module doing the thing.",
    )
    parser.add_argument(
        "--module_do_the_thing",
        type=bool,
        default=True,
        help="Should the module do the thing?"
    )

    main.py:
    --------
    import config

    def main() -> None:
        parser = config.add_commandline_args(
            "Main",
            "A program that does the thing.",
        )
        parser.add_argument(
            "--dry_run",
            type=bool,
            default=False,
            help="Should we really do the thing?"
        )
        config.parse()   # Very important, this must be invoked!

    If you set this up and remember to invoke config.parse(), all commandline
    arguments will play nicely together:

    % main.py -h
    usage: main.py [-h]
                   [--module_do_the_thing MODULE_DO_THE_THING]
                   [--dry_run DRY_RUN]

    Module:
      Args related to module doing the thing.

      --module_do_the_thing MODULE_DO_THE_THING
                   Should the module do the thing?

    Main:
      A program that does the thing

      --dry_run
                   Should we really do the thing?

    Arguments themselves should be accessed via
    config.config['arg_name'].  e.g.

    if not config.config['dry_run']:
        module.do_the_thing()

"""

import argparse
import logging
import os
import pprint
import re
import sys
from typing import Any, Dict, List

# Note: at this point in time, logging hasn't been configured and
# anything we log will come out the root logger.


class LoadFromFile(argparse.Action):
    """Helper to load a config file into argparse."""
    def __call__ (self, parser, namespace, values, option_string = None):
        with values as f:
            buf = f.read()
            argv = []
            for line in buf.split(','):
                line = line.strip()
                line = line.strip('{')
                line = line.strip('}')
                m = re.match(r"^'([a-zA-Z_\-]+)'\s*:\s*(.*)$", line)
                if m:
                    key = m.group(1)
                    value = m.group(2)
                    value = value.strip("'")
                    if value not in ('None', 'True', 'False'):
                        argv.append(f'--{key}')
                        argv.append(value)
            parser.parse_args(argv, namespace)


# A global parser that we will collect arguments into.
args = argparse.ArgumentParser(
    description=f"This program uses config.py ({__file__}) for global, cross-module configuration.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    fromfile_prefix_chars="@",
)
config_parse_called = False

# A global configuration dictionary that will contain parsed arguments
# It is also this variable that modules use to access parsed arguments
config: Dict[str, Any] = {}

# Defer logging messages until later when logging has been initialized.
saved_messages: List[str] = []


def add_commandline_args(title: str, description: str = ""):
    """Create a new context for arguments and return a handle."""
    return args.add_argument_group(title, description)


group = add_commandline_args(
    f'Global Config ({__file__})',
    'Args that control the global config itself; how meta!',
)
group.add_argument(
    '--config_loadfile',
    type=open,
    action=LoadFromFile,
    metavar='FILENAME',
    default=None,
    help='Config file from which to read args in lieu or in addition to commandline.',
)
group.add_argument(
    '--config_dump',
    default=False,
    action='store_true',
    help='Display the global configuration on STDERR at program startup.',
)
group.add_argument(
    '--config_savefile',
    type=str,
    metavar='FILENAME',
    default=None,
    help='Populate config file compatible --config_loadfile to save config for later use.',
)


def is_flag_already_in_argv(var: str):
    """Is a particular flag passed on the commandline?"""
    for _ in sys.argv:
        if var in _:
            return True
    return False


def parse() -> Dict[str, Any]:
    import string_utils

    """Main program should call this early in main()"""
    global config_parse_called
    if config_parse_called:
        return
    config_parse_called = True
    global saved_messages

    # Examine the environment variables to settings that match
    # known flags.
    usage_message = args.format_usage()
    optional = False
    var = ''
    for x in usage_message.split():
        if x[0] == '[':
            optional = True
        if optional:
            var += f'{x} '
            if x[-1] == ']':
                optional = False
                var = var.strip()
                var = var.strip('[')
                var = var.strip(']')
                chunks = var.split()
                if len(chunks) > 1:
                    var = var.split()[0]

                # Environment vars the same as flag names without
                # the initial -'s and in UPPERCASE.
                env = var.strip('-').upper()
                if env in os.environ:
                    if not is_flag_already_in_argv(var):
                        value = os.environ[env]
                        saved_messages.append(
                            f'Initialized from environment: {var} = {value}'
                        )
                        if len(chunks) == 1 and string_utils.to_bool(value):
                            sys.argv.append(var)
                        elif len(chunks) > 1:
                            sys.argv.append(var)
                            sys.argv.append(value)
                var = ''
                env = ''
        else:
            next

    # Parse (possibly augmented) commandline args with argparse normally.
    #config.update(vars(args.parse_args()))
    known, unknown = args.parse_known_args()
    config.update(vars(known))

    # Reconstruct the argv with unrecognized flags for the benefit of
    # future argument parsers.
    sys.argv = sys.argv[:1] + unknown

    if config['config_savefile']:
        with open(config['config_savefile'], 'w') as wf:
            wf.write("\n".join(sys.argv[1:]))

    if config['config_dump']:
        dump_config()
    return config


def has_been_parsed() -> bool:
    """Has the global config been parsed yet?"""
    global config_parse_called
    return config_parse_called


def dump_config():
    """Print the current config to stdout."""
    print("Global Configuration:", file=sys.stderr)
    pprint.pprint(config, stream=sys.stderr)


def late_logging():
    """Log messages saved earlier now that logging has been initialized."""
    logger = logging.getLogger(__name__)
    global saved_messages
    for _ in saved_messages:
        logger.debug(_)