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