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 class OptionalRawFormatter(argparse.HelpFormatter):
  88     """This formatter has the same bahavior as the normal argparse text
  89     formatter except when the help text of an argument begins with
  90     "RAW|".  In that case, the line breaks are preserved and the text
  91     is not wrapped.
  92
  93     """
  94
  95     def _split_lines(self, text, width):
  96         if text.startswith('RAW|'):
  97             return text[4:].splitlines()
  98         return argparse.HelpFormatter._split_lines(self, text, width)
  99
 100
 101 # A global parser that we will collect arguments into.
 102 args = argparse.ArgumentParser(
 103     description=None,
 104     formatter_class=OptionalRawFormatter,
 105     fromfile_prefix_chars="@",
 106     epilog=f'{program_name} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
 107 )
 108
 109 # Keep track of if we've been called and prevent being called more
 110 # than once.
 111 config_parse_called = False
 112
 113
 114 # A global configuration dictionary that will contain parsed arguments.
 115 # It is also this variable that modules use to access parsed arguments.
 116 # This is the data that is most interesting to our callers; it will hold
 117 # the configuration result.
 118 config: Dict[str, Any] = {}
 119 # It would be really nice if this shit worked from interactive python
 120
 121
 122 def add_commandline_args(title: str, description: str = ""):
 123     """Create a new context for arguments and return a handle."""
 124     return args.add_argument_group(title, description)
 125
 126
 127 group = add_commandline_args(
 128     f'Global Config ({__file__})',
 129     'Args that control the global config itself; how meta!',
 130 )
 131 group.add_argument(
 132     '--config_loadfile',
 133     metavar='FILENAME',
 134     default=None,
 135     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 136 )
 137 group.add_argument(
 138     '--config_dump',
 139     default=False,
 140     action='store_true',
 141     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 142 )
 143 group.add_argument(
 144     '--config_savefile',
 145     type=str,
 146     metavar='FILENAME',
 147     default=None,
 148     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 149 )
 150 group.add_argument(
 151     '--config_rejects_unrecognized_arguments',
 152     default=False,
 153     action='store_true',
 154     help=(
 155         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 156         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 157         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 158     ),
 159 )
 160
 161
 162 def overwrite_argparse_epilog(msg: str) -> None:
 163     args.epilog = msg
 164
 165
 166 def is_flag_already_in_argv(var: str):
 167     """Is a particular flag passed on the commandline?"""
 168     for _ in sys.argv:
 169         if var in _:
 170             return True
 171     return False
 172
 173
 174 def reorder_arg_action_groups(entry_module: Optional[str]):
 175     global program_name, args
 176     reordered_action_groups = []
 177     for group in args._action_groups:
 178         if entry_module is not None and entry_module in group.title:  # type: ignore
 179             reordered_action_groups.append(group)
 180         elif program_name in group.title:  # type: ignore
 181             reordered_action_groups.append(group)
 182         else:
 183             reordered_action_groups.insert(0, group)
 184     return reordered_action_groups
 185
 186
 187 def augment_sys_argv_from_environment_variables():
 188     global saved_messages
 189     usage_message = args.format_usage()
 190     optional = False
 191     var = ''
 192     for x in usage_message.split():
 193         if x[0] == '[':
 194             optional = True
 195         if optional:
 196             var += f'{x} '
 197             if x[-1] == ']':
 198                 optional = False
 199                 var = var.strip()
 200                 var = var.strip('[')
 201                 var = var.strip(']')
 202                 chunks = var.split()
 203                 if len(chunks) > 1:
 204                     var = var.split()[0]
 205
 206                 # Environment vars the same as flag names without
 207                 # the initial -'s and in UPPERCASE.
 208                 env = var.strip('-').upper()
 209                 if env in os.environ:
 210                     if not is_flag_already_in_argv(var):
 211                         value = os.environ[env]
 212                         saved_messages.append(f'Initialized from environment: {var} = {value}')
 213                         from string_utils import to_bool
 214
 215                         if len(chunks) == 1 and to_bool(value):
 216                             sys.argv.append(var)
 217                         elif len(chunks) > 1:
 218                             sys.argv.append(var)
 219                             sys.argv.append(value)
 220                 var = ''
 221                 env = ''
 222
 223
 224 def augment_sys_argv_from_loadfile():
 225     global saved_messages
 226     loadfile = None
 227     saw_other_args = False
 228     grab_next_arg = False
 229     for arg in sys.argv[1:]:
 230         if 'config_loadfile' in arg:
 231             pieces = arg.split('=')
 232             if len(pieces) > 1:
 233                 loadfile = pieces[1]
 234             else:
 235                 grab_next_arg = True
 236         elif grab_next_arg:
 237             loadfile = arg
 238         else:
 239             saw_other_args = True
 240
 241     if loadfile is not None:
 242         if not os.path.exists(loadfile):
 243             raise Exception(
 244                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 245             )
 246         if saw_other_args:
 247             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 248         else:
 249             msg = f'Reading commandline arguments from {loadfile}.'
 250         print(msg, file=sys.stderr)
 251         saved_messages.append(msg)
 252
 253         with open(loadfile, 'r') as rf:
 254             newargs = rf.readlines()
 255         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 256         sys.argv += newargs
 257
 258
 259 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 260     """Main program should call this early in main().  Note that the
 261     bootstrap.initialize wrapper takes care of this automatically.
 262
 263     """
 264     global config_parse_called
 265     if config_parse_called:
 266         return config
 267     global saved_messages
 268
 269     # If we're about to do the usage message dump, put the main
 270     # module's argument group last in the list (if possible) so that
 271     # when the user passes -h or --help, it will be visible on the
 272     # screen w/o scrolling.
 273     for arg in sys.argv:
 274         if arg == '--help' or arg == '-h':
 275             args._action_groups = reorder_arg_action_groups(entry_module)
 276
 277     # Examine the environment for variables that match known flags.
 278     # For a flag called --example_flag the corresponding environment
 279     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 280     # these into sys.argv to be parsed.
 281     augment_sys_argv_from_environment_variables()
 282
 283     # Look for loadfile and read/parse it if present.  This also
 284     # works by jamming these values onto sys.argv.
 285     augment_sys_argv_from_loadfile()
 286
 287     # Parse (possibly augmented, possibly completely overwritten)
 288     # commandline args with argparse normally and populate config.
 289     known, unknown = args.parse_known_args()
 290     config.update(vars(known))
 291
 292     # Reconstruct the argv with unrecognized flags for the benefit of
 293     # future argument parsers.  For example, unittest_main in python
 294     # has some of its own flags.  If we didn't recognize it, maybe
 295     # someone else will.
 296     if len(unknown) > 0:
 297         if config['config_rejects_unrecognized_arguments']:
 298             raise Exception(
 299                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 300             )
 301         saved_messages.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 302     sys.argv = sys.argv[:1] + unknown
 303
 304     # Check for savefile and populate it if requested.
 305     savefile = config['config_savefile']
 306     if savefile and len(savefile) > 0:
 307         with open(savefile, 'w') as wf:
 308             wf.write("\n".join(original_argv[1:]))
 309
 310     # Also dump the config on stderr if requested.
 311     if config['config_dump']:
 312         dump_config()
 313
 314     config_parse_called = True
 315     return config
 316
 317
 318 def has_been_parsed() -> bool:
 319     """Has the global config been parsed yet?"""
 320     global config_parse_called
 321     return config_parse_called
 322
 323
 324 def dump_config():
 325     """Print the current config to stdout."""
 326     print("Global Configuration:", file=sys.stderr)
 327     pprint.pprint(config, stream=sys.stderr)
 328     print()
 329
 330
 331 def late_logging():
 332     """Log messages saved earlier now that logging has been initialized."""
 333     logger = logging.getLogger(__name__)
 334     global saved_messages
 335     for _ in saved_messages:
 336         logger.debug(_)