path: root/config.py

#!/usr/bin/env python3

"""Global configuration driven by commandline arguments, environment variables
and saved configuration files.  This works across several 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.  This is done automatically for you
    if you're using the bootstrap module's initialize wrapper.

    % 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 sys
from typing import Any, Dict, List, Optional

# This module is commonly used by others in here and should avoid
# taking any unnecessary dependencies back on them.

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

# Make a copy of the original program arguments.
program_name = os.path.basename(sys.argv[0])
original_argv = [arg for arg in sys.argv]


# A global parser that we will collect arguments into.
args = argparse.ArgumentParser(
    description=None,
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    fromfile_prefix_chars="@",
    epilog=f'{program_name} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.'
)

# Keep track of if we've been called and prevent being called more
# than once.
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.
# This is the data that is most interesting to our callers; it will hold
# the configuration result.
config = {}
# It would be really nice if this shit worked from interactive python


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',
    metavar='FILENAME',
    default=None,
    help='Config file (populated via --config_savefile) 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 (possibly derived from multiple sources) on STDERR at program startup.',
)
group.add_argument(
    '--config_savefile',
    type=str,
    metavar='FILENAME',
    default=None,
    help='Populate config file compatible with --config_loadfile to save global 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(entry_module: Optional[str]) -> Dict[str, Any]:
    """Main program should call this early in main().  Note that the
    bootstrap.initialize wrapper takes care of this automatically.

    """
    global config_parse_called
    if config_parse_called:
        return config

    global saved_messages

    # If we're about to do the usage message dump, put the main module's
    # argument group last in the list (if possible) so that when the user
    # passes -h or --help, it will be visible on the screen w/o scrolling.
    reordered_action_groups = []
    global prog
    for arg in sys.argv:
        if arg == '--help' or arg == '-h':
            for group in args._action_groups:
                if entry_module is not None and entry_module in group.title:
                    reordered_action_groups.append(group)
                elif program_name in group.title:
                    reordered_action_groups.append(group)
                else:
                    reordered_action_groups.insert(0, group)
            args._action_groups = reordered_action_groups

    # Examine the environment for variables that match known flags.
    # For a flag called --example_flag the corresponding environment
    # variable would be called EXAMPLE_FLAG.
    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}'
                        )
                        from string_utils import to_bool
                        if len(chunks) == 1 and to_bool(value):
                            sys.argv.append(var)
                        elif len(chunks) > 1:
                            sys.argv.append(var)
                            sys.argv.append(value)
                var = ''
                env = ''
        else:
            next

    # Look for loadfile and read/parse it if present.
    loadfile = None
    saw_other_args = False
    grab_next_arg = False
    for arg in sys.argv[1:]:
        if 'config_loadfile' in arg:
            pieces = arg.split('=')
            if len(pieces) > 1:
                loadfile = pieces[1]
            else:
                grab_next_arg = True
        elif grab_next_arg:
            loadfile = arg
        else:
            saw_other_args = True

    if loadfile is not None:
        if saw_other_args:
            msg = f'Augmenting commandline arguments with those from {loadfile}.'
            print(msg, file=sys.stderr)
            saved_messages.append(msg)
        if not os.path.exists(loadfile):
            print(f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.',
                  file=sys.stderr)
            sys.exit(-1)
        with open(loadfile, 'r') as rf:
            newargs = rf.readlines()
        newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
        sys.argv += newargs

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

    # Reconstruct the argv with unrecognized flags for the benefit of
    # future argument parsers.  For example, unittest_main in python
    # has some of its own flags.  If we didn't recognize it, maybe
    # someone else will.
    sys.argv = sys.argv[:1] + unknown

    # Check for savefile and populate it if requested.
    savefile = config['config_savefile']
    if savefile and len(savefile) > 0:
        with open(savefile, 'w') as wf:
            wf.write(
                "\n".join(original_argv[1:])
            )

    # Also dump the config on stderr if requested.
    if config['config_dump']:
        dump_config()

    config_parse_called = True
    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)
    print()


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(_)