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 = os.path.basename(sys.argv[0])
  84 original_argv = [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 = {}
 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     reordered_action_groups = []
 158     for group in args._action_groups:
 159         if entry_module is not None and entry_module in group.title:
 160             reordered_action_groups.append(group)
 161         elif program_name in group.title:
 162             reordered_action_groups.append(group)
 163         else:
 164             reordered_action_groups.insert(0, group)
 165     return reordered_action_groups
 166
 167
 168 def augment_sys_argv_from_environment_variables():
 169     global saved_messages
 170     usage_message = args.format_usage()
 171     optional = False
 172     var = ''
 173     for x in usage_message.split():
 174         if x[0] == '[':
 175             optional = True
 176         if optional:
 177             var += f'{x} '
 178             if x[-1] == ']':
 179                 optional = False
 180                 var = var.strip()
 181                 var = var.strip('[')
 182                 var = var.strip(']')
 183                 chunks = var.split()
 184                 if len(chunks) > 1:
 185                     var = var.split()[0]
 186
 187                 # Environment vars the same as flag names without
 188                 # the initial -'s and in UPPERCASE.
 189                 env = var.strip('-').upper()
 190                 if env in os.environ:
 191                     if not is_flag_already_in_argv(var):
 192                         value = os.environ[env]
 193                         saved_messages.append(
 194                             f'Initialized from environment: {var} = {value}'
 195                         )
 196                         from string_utils import to_bool
 197
 198                         if len(chunks) == 1 and to_bool(value):
 199                             sys.argv.append(var)
 200                         elif len(chunks) > 1:
 201                             sys.argv.append(var)
 202                             sys.argv.append(value)
 203                 var = ''
 204                 env = ''
 205
 206
 207 def augment_sys_argv_from_loadfile():
 208     global saved_messages
 209     loadfile = None
 210     saw_other_args = False
 211     grab_next_arg = False
 212     for arg in sys.argv[1:]:
 213         if 'config_loadfile' in arg:
 214             pieces = arg.split('=')
 215             if len(pieces) > 1:
 216                 loadfile = pieces[1]
 217             else:
 218                 grab_next_arg = True
 219         elif grab_next_arg:
 220             loadfile = arg
 221         else:
 222             saw_other_args = True
 223
 224     if loadfile is not None:
 225         if not os.path.exists(loadfile):
 226             raise Exception(
 227                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 228             )
 229         if saw_other_args:
 230             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 231         else:
 232             msg = f'Reading commandline arguments from {loadfile}.'
 233         print(msg, file=sys.stderr)
 234         saved_messages.append(msg)
 235
 236         with open(loadfile, 'r') as rf:
 237             newargs = rf.readlines()
 238         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 239         sys.argv += newargs
 240
 241
 242 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 243     """Main program should call this early in main().  Note that the
 244     bootstrap.initialize wrapper takes care of this automatically.
 245
 246     """
 247     global config_parse_called
 248     if config_parse_called:
 249         return config
 250     global saved_messages
 251
 252     # If we're about to do the usage message dump, put the main
 253     # module's argument group last in the list (if possible) so that
 254     # when the user passes -h or --help, it will be visible on the
 255     # screen w/o scrolling.
 256     for arg in sys.argv:
 257         if arg == '--help' or arg == '-h':
 258             args._action_groups = reorder_arg_action_groups(entry_module)
 259
 260     # Examine the environment for variables that match known flags.
 261     # For a flag called --example_flag the corresponding environment
 262     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 263     # these into sys.argv to be parsed.
 264     augment_sys_argv_from_environment_variables()
 265
 266     # Look for loadfile and read/parse it if present.  This also
 267     # works by jamming these values onto sys.argv.
 268     augment_sys_argv_from_loadfile()
 269
 270     # Parse (possibly augmented, possibly completely overwritten)
 271     # commandline args with argparse normally and populate config.
 272     known, unknown = args.parse_known_args()
 273     config.update(vars(known))
 274
 275     # Reconstruct the argv with unrecognized flags for the benefit of
 276     # future argument parsers.  For example, unittest_main in python
 277     # has some of its own flags.  If we didn't recognize it, maybe
 278     # someone else will.
 279     if len(unknown) > 0:
 280         if config['config_rejects_unrecognized_arguments']:
 281             raise Exception(
 282                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 283             )
 284         saved_messages.append(
 285             f'Config encountered unrecognized commandline arguments: {unknown}'
 286         )
 287     sys.argv = sys.argv[:1] + unknown
 288
 289     # Check for savefile and populate it if requested.
 290     savefile = config['config_savefile']
 291     if savefile and len(savefile) > 0:
 292         with open(savefile, 'w') as wf:
 293             wf.write("\n".join(original_argv[1:]))
 294
 295     # Also dump the config on stderr if requested.
 296     if config['config_dump']:
 297         dump_config()
 298
 299     config_parse_called = True
 300     return config
 301
 302
 303 def has_been_parsed() -> bool:
 304     """Has the global config been parsed yet?"""
 305     global config_parse_called
 306     return config_parse_called
 307
 308
 309 def dump_config():
 310     """Print the current config to stdout."""
 311     print("Global Configuration:", file=sys.stderr)
 312     pprint.pprint(config, stream=sys.stderr)
 313     print()
 314
 315
 316 def late_logging():
 317     """Log messages saved earlier now that logging has been initialized."""
 318     logger = logging.getLogger(__name__)
 319     global saved_messages
 320     for _ in saved_messages:
 321         logger.debug(_)