config.py

   1 #!/usr/bin/env python3
   2
   3 """Global configuration driven by commandline arguments, environment variables
   4 and saved configuration files.  This works across several modules.
   5
   6 Usage:
   7
   8     module.py:
   9     ----------
  10     import config
  11
  12     parser = config.add_commandline_args(
  13         "Module",
  14         "Args related to module doing the thing.",
  15     )
  16     parser.add_argument(
  17         "--module_do_the_thing",
  18         type=bool,
  19         default=True,
  20         help="Should the module do the thing?"
  21     )
  22
  23     main.py:
  24     --------
  25     import config
  26
  27     def main() -> None:
  28         parser = config.add_commandline_args(
  29             "Main",
  30             "A program that does the thing.",
  31         )
  32         parser.add_argument(
  33             "--dry_run",
  34             type=bool,
  35             default=False,
  36             help="Should we really do the thing?"
  37         )
  38         config.parse()   # Very important, this must be invoked!
  39
  40     If you set this up and remember to invoke config.parse(), all commandline
  41     arguments will play nicely together.  This is done automatically for you
  42     if you're using the bootstrap module's initialize wrapper.
  43
  44     % main.py -h
  45     usage: main.py [-h]
  46                    [--module_do_the_thing MODULE_DO_THE_THING]
  47                    [--dry_run DRY_RUN]
  48
  49     Module:
  50       Args related to module doing the thing.
  51
  52       --module_do_the_thing MODULE_DO_THE_THING
  53                    Should the module do the thing?
  54
  55     Main:
  56       A program that does the thing
  57
  58       --dry_run
  59                    Should we really do the thing?
  60
  61     Arguments themselves should be accessed via
  62     config.config['arg_name'].  e.g.
  63
  64     if not config.config['dry_run']:
  65         module.do_the_thing()
  66
  67 """
  68
  69 import argparse
  70 import logging
  71 import os
  72 import pprint
  73 import sys
  74 from typing import Any, Dict, List, Optional
  75
  76 # This module is commonly used by others in here and should avoid
  77 # taking any unnecessary dependencies back on them.
  78
  79 # Defer logging messages until later when logging has been initialized.
  80 saved_messages: List[str] = []
  81
  82 # Make a copy of the original program arguments.
  83 program_name: str = os.path.basename(sys.argv[0])
  84 original_argv: List[str] = [arg for arg in sys.argv]
  85
  86
  87 # A global parser that we will collect arguments into.
  88 args = argparse.ArgumentParser(
  89     description=None,
  90     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  91     fromfile_prefix_chars="@",
  92     epilog=f'{program_name} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
  93 )
  94
  95 # Keep track of if we've been called and prevent being called more
  96 # than once.
  97 config_parse_called = False
  98
  99
 100 # A global configuration dictionary that will contain parsed arguments.
 101 # It is also this variable that modules use to access parsed arguments.
 102 # This is the data that is most interesting to our callers; it will hold
 103 # the configuration result.
 104 config: Dict[str, Any] = {}
 105 # It would be really nice if this shit worked from interactive python
 106
 107
 108 def add_commandline_args(title: str, description: str = ""):
 109     """Create a new context for arguments and return a handle."""
 110     return args.add_argument_group(title, description)
 111
 112
 113 group = add_commandline_args(
 114     f'Global Config ({__file__})',
 115     'Args that control the global config itself; how meta!',
 116 )
 117 group.add_argument(
 118     '--config_loadfile',
 119     metavar='FILENAME',
 120     default=None,
 121     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 122 )
 123 group.add_argument(
 124     '--config_dump',
 125     default=False,
 126     action='store_true',
 127     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 128 )
 129 group.add_argument(
 130     '--config_savefile',
 131     type=str,
 132     metavar='FILENAME',
 133     default=None,
 134     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 135 )
 136 group.add_argument(
 137     '--config_rejects_unrecognized_arguments',
 138     default=False,
 139     action='store_true',
 140     help=(
 141         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 142         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 143         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 144     ),
 145 )
 146
 147
 148 def is_flag_already_in_argv(var: str):
 149     """Is a particular flag passed on the commandline?"""
 150     for _ in sys.argv:
 151         if var in _:
 152             return True
 153     return False
 154
 155
 156 def reorder_arg_action_groups(entry_module: Optional[str]):
 157     global program_name, args
 158     reordered_action_groups = []
 159     for group in args._action_groups:
 160         if entry_module is not None and entry_module in group.title:  # type: ignore
 161             reordered_action_groups.append(group)
 162         elif program_name in group.title:  # type: ignore
 163             reordered_action_groups.append(group)
 164         else:
 165             reordered_action_groups.insert(0, group)
 166     return reordered_action_groups
 167
 168
 169 def augment_sys_argv_from_environment_variables():
 170     global saved_messages
 171     usage_message = args.format_usage()
 172     optional = False
 173     var = ''
 174     for x in usage_message.split():
 175         if x[0] == '[':
 176             optional = True
 177         if optional:
 178             var += f'{x} '
 179             if x[-1] == ']':
 180                 optional = False
 181                 var = var.strip()
 182                 var = var.strip('[')
 183                 var = var.strip(']')
 184                 chunks = var.split()
 185                 if len(chunks) > 1:
 186                     var = var.split()[0]
 187
 188                 # Environment vars the same as flag names without
 189                 # the initial -'s and in UPPERCASE.
 190                 env = var.strip('-').upper()
 191                 if env in os.environ:
 192                     if not is_flag_already_in_argv(var):
 193                         value = os.environ[env]
 194                         saved_messages.append(
 195                             f'Initialized from environment: {var} = {value}'
 196                         )
 197                         from string_utils import to_bool
 198
 199                         if len(chunks) == 1 and to_bool(value):
 200                             sys.argv.append(var)
 201                         elif len(chunks) > 1:
 202                             sys.argv.append(var)
 203                             sys.argv.append(value)
 204                 var = ''
 205                 env = ''
 206
 207
 208 def augment_sys_argv_from_loadfile():
 209     global saved_messages
 210     loadfile = None
 211     saw_other_args = False
 212     grab_next_arg = False
 213     for arg in sys.argv[1:]:
 214         if 'config_loadfile' in arg:
 215             pieces = arg.split('=')
 216             if len(pieces) > 1:
 217                 loadfile = pieces[1]
 218             else:
 219                 grab_next_arg = True
 220         elif grab_next_arg:
 221             loadfile = arg
 222         else:
 223             saw_other_args = True
 224
 225     if loadfile is not None:
 226         if not os.path.exists(loadfile):
 227             raise Exception(
 228                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 229             )
 230         if saw_other_args:
 231             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 232         else:
 233             msg = f'Reading commandline arguments from {loadfile}.'
 234         print(msg, file=sys.stderr)
 235         saved_messages.append(msg)
 236
 237         with open(loadfile, 'r') as rf:
 238             newargs = rf.readlines()
 239         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 240         sys.argv += newargs
 241
 242
 243 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 244     """Main program should call this early in main().  Note that the
 245     bootstrap.initialize wrapper takes care of this automatically.
 246
 247     """
 248     global config_parse_called
 249     if config_parse_called:
 250         return config
 251     global saved_messages
 252
 253     # If we're about to do the usage message dump, put the main
 254     # module's argument group last in the list (if possible) so that
 255     # when the user passes -h or --help, it will be visible on the
 256     # screen w/o scrolling.
 257     for arg in sys.argv:
 258         if arg == '--help' or arg == '-h':
 259             args._action_groups = reorder_arg_action_groups(entry_module)
 260
 261     # Examine the environment for variables that match known flags.
 262     # For a flag called --example_flag the corresponding environment
 263     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 264     # these into sys.argv to be parsed.
 265     augment_sys_argv_from_environment_variables()
 266
 267     # Look for loadfile and read/parse it if present.  This also
 268     # works by jamming these values onto sys.argv.
 269     augment_sys_argv_from_loadfile()
 270
 271     # Parse (possibly augmented, possibly completely overwritten)
 272     # commandline args with argparse normally and populate config.
 273     known, unknown = args.parse_known_args()
 274     config.update(vars(known))
 275
 276     # Reconstruct the argv with unrecognized flags for the benefit of
 277     # future argument parsers.  For example, unittest_main in python
 278     # has some of its own flags.  If we didn't recognize it, maybe
 279     # someone else will.
 280     if len(unknown) > 0:
 281         if config['config_rejects_unrecognized_arguments']:
 282             raise Exception(
 283                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 284             )
 285         saved_messages.append(
 286             f'Config encountered unrecognized commandline arguments: {unknown}'
 287         )
 288     sys.argv = sys.argv[:1] + unknown
 289
 290     # Check for savefile and populate it if requested.
 291     savefile = config['config_savefile']
 292     if savefile and len(savefile) > 0:
 293         with open(savefile, 'w') as wf:
 294             wf.write("\n".join(original_argv[1:]))
 295
 296     # Also dump the config on stderr if requested.
 297     if config['config_dump']:
 298         dump_config()
 299
 300     config_parse_called = True
 301     return config
 302
 303
 304 def has_been_parsed() -> bool:
 305     """Has the global config been parsed yet?"""
 306     global config_parse_called
 307     return config_parse_called
 308
 309
 310 def dump_config():
 311     """Print the current config to stdout."""
 312     print("Global Configuration:", file=sys.stderr)
 313     pprint.pprint(config, stream=sys.stderr)
 314     print()
 315
 316
 317 def late_logging():
 318     """Log messages saved earlier now that logging has been initialized."""
 319     logger = logging.getLogger(__name__)
 320     global saved_messages
 321     for _ in saved_messages:
 322         logger.debug(_)