config.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Global configuration driven by commandline arguments, environment variables
   6 and saved configuration files.  This works across several modules.
   7
   8 Usage:
   9
  10     In your file.py::
  11
  12         import config
  13
  14         parser = config.add_commandline_args(
  15             "Module",
  16             "Args related to module doing the thing.",
  17         )
  18         parser.add_argument(
  19             "--module_do_the_thing",
  20             type=bool,
  21             default=True,
  22             help="Should the module do the thing?"
  23         )
  24
  25     In your main.py::
  26
  27         import config
  28
  29         def main() -> None:
  30             parser = config.add_commandline_args(
  31                 "Main",
  32                 "A program that does the thing.",
  33             )
  34             parser.add_argument(
  35                 "--dry_run",
  36                 type=bool,
  37                 default=False,
  38                 help="Should we really do the thing?"
  39             )
  40             config.parse()   # Very important, this must be invoked!
  41
  42     If you set this up and remember to invoke config.parse(), all commandline
  43     arguments will play nicely together.  This is done automatically for you
  44     if you're using the :meth:`bootstrap.initialize` decorator on
  45     your program's entry point.  See :meth:`python_modules.bootstrap.initialize`
  46     for more details.::
  47
  48         import bootstrap
  49
  50         @bootstrap.initialize
  51         def main():
  52             whatever
  53
  54         if __name__ == '__main__':
  55             main()
  56
  57     Either way, you'll get this behavior from the commandline::
  58
  59         % main.py -h
  60         usage: main.py [-h]
  61                        [--module_do_the_thing MODULE_DO_THE_THING]
  62                        [--dry_run DRY_RUN]
  63
  64         Module:
  65           Args related to module doing the thing.
  66
  67           --module_do_the_thing MODULE_DO_THE_THING
  68                        Should the module do the thing?
  69
  70         Main:
  71           A program that does the thing
  72
  73           --dry_run
  74                        Should we really do the thing?
  75
  76     Arguments themselves should be accessed via
  77     :code:`config.config['arg_name']`.  e.g.::
  78
  79         if not config.config['dry_run']:
  80             module.do_the_thing()
  81
  82 """
  83
  84 import argparse
  85 import logging
  86 import os
  87 import pprint
  88 import sys
  89 from typing import Any, Dict, List, Optional
  90
  91 # This module is commonly used by others in here and should avoid
  92 # taking any unnecessary dependencies back on them.
  93
  94 # Defer logging messages until later when logging has been initialized.
  95 SAVED_MESSAGES: List[str] = []
  96
  97 # Make a copy of the original program arguments.
  98 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
  99 ORIG_ARGV: List[str] = sys.argv.copy()
 100
 101
 102 class OptionalRawFormatter(argparse.HelpFormatter):
 103     """This formatter has the same bahavior as the normal argparse text
 104     formatter except when the help text of an argument begins with
 105     "RAW|".  In that case, the line breaks are preserved and the text
 106     is not wrapped.
 107
 108     Use this, for example, when you need the helptext of an argument
 109     to have its spacing preserved exactly, e.g.::
 110
 111         args.add_argument(
 112             '--mode',
 113             type=str,
 114             default='PLAY',
 115             choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
 116             metavar='MODE',
 117             help='''RAW|Our mode of operation.  One of:
 118
 119                 PLAY = play wordle with me!  Pick a random solution or
 120                        specify a solution with --template.
 121
 122                CHEAT = given a --template and, optionally, --letters_in_word
 123                        and/or --letters_to_avoid, return the best guess word;
 124
 125             AUTOPLAY = given a complete word in --template, guess it step
 126                        by step showing work;
 127
 128             SELFTEST = autoplay every possible solution keeping track of
 129                        wins/losses and average number of guesses;
 130
 131           PRECOMPUTE = populate hash table with optimal guesses.
 132             ''',
 133         )
 134
 135     """
 136
 137     def _split_lines(self, text, width):
 138         if text.startswith('RAW|'):
 139             return text[4:].splitlines()
 140         return argparse.HelpFormatter._split_lines(self, text, width)
 141
 142
 143 # A global parser that we will collect arguments into.
 144 ARGS = argparse.ArgumentParser(
 145     description=None,
 146     formatter_class=OptionalRawFormatter,
 147     fromfile_prefix_chars="@",
 148     epilog=f'{PROGRAM_NAME} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
 149     # I don't fully understand why but when loaded by sphinx sometimes
 150     # the same module is loaded many times causing any arguments it
 151     # registers via module-level code to be redefined.  Work around
 152     # this iff the program is 'sphinx-build'
 153     conflict_handler='resolve' if PROGRAM_NAME == 'sphinx-build' else 'error',
 154 )
 155
 156 # Keep track of if we've been called and prevent being called more
 157 # than once.
 158 CONFIG_PARSE_CALLED = False
 159
 160
 161 # A global configuration dictionary that will contain parsed arguments.
 162 # It is also this variable that modules use to access parsed arguments.
 163 # This is the data that is most interesting to our callers; it will hold
 164 # the configuration result.
 165 config: Dict[str, Any] = {}
 166
 167 # It would be really nice if this shit worked from interactive python
 168
 169
 170 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
 171     """Create a new context for arguments and return a handle.
 172
 173     Args:
 174         title: A title for your module's commandline arguments group.
 175         description: A helpful description of your module.
 176
 177     Returns:
 178         An argparse._ArgumentGroup to be populated by the caller.
 179     """
 180     return ARGS.add_argument_group(title, description)
 181
 182
 183 group = add_commandline_args(
 184     f'Global Config ({__file__})',
 185     'Args that control the global config itself; how meta!',
 186 )
 187 group.add_argument(
 188     '--config_loadfile',
 189     metavar='FILENAME',
 190     default=None,
 191     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 192 )
 193 group.add_argument(
 194     '--config_dump',
 195     default=False,
 196     action='store_true',
 197     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 198 )
 199 group.add_argument(
 200     '--config_savefile',
 201     type=str,
 202     metavar='FILENAME',
 203     default=None,
 204     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 205 )
 206 group.add_argument(
 207     '--config_rejects_unrecognized_arguments',
 208     default=False,
 209     action='store_true',
 210     help=(
 211         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 212         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 213         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 214     ),
 215 )
 216
 217
 218 def overwrite_argparse_epilog(msg: str) -> None:
 219     """Allows your code to override the default epilog created by
 220     argparse.
 221
 222     Args:
 223         msg: The epilog message to substitute for the default.
 224     """
 225     ARGS.epilog = msg
 226
 227
 228 def is_flag_already_in_argv(var: str) -> bool:
 229     """Returns true if a particular flag is passed on the commandline?
 230
 231     Args:
 232         var: The flag to search for.
 233     """
 234     for _ in sys.argv:
 235         if var in _:
 236             return True
 237     return False
 238
 239
 240 def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
 241     """Internal.  Used to reorder the arguments before dumping out a
 242     generated help string such that the main program's arguments come
 243     last.
 244
 245     """
 246     reordered_action_groups = []
 247     for grp in ARGS._action_groups:
 248         if entry_module is not None and entry_module in grp.title:  # type: ignore
 249             reordered_action_groups.append(grp)
 250         elif PROGRAM_NAME in group.title:  # type: ignore
 251             reordered_action_groups.append(grp)
 252         else:
 253             reordered_action_groups.insert(0, grp)
 254     return reordered_action_groups
 255
 256
 257 def print_usage() -> None:
 258     """Prints the normal help usage message out."""
 259     ARGS.print_help()
 260
 261
 262 def usage() -> str:
 263     """
 264     Returns:
 265         program usage help text as a string.
 266     """
 267     return ARGS.format_usage()
 268
 269
 270 def _augment_sys_argv_from_environment_variables():
 271     """Internal.  Look at the system environment for variables that match
 272     arg names.  This is done via some munging such that:
 273
 274     :code:`--argument_to_match`
 275
 276     ...is matched by:
 277
 278     :code:`ARGUMENT_TO_MATCH`
 279
 280     This allows programmers to set args via shell environment variables
 281     in lieu of passing them on the cmdline.
 282
 283     """
 284
 285     usage_message = usage()
 286     optional = False
 287     var = ''
 288     for x in usage_message.split():
 289         if x[0] == '[':
 290             optional = True
 291         if optional:
 292             var += f'{x} '
 293             if x[-1] == ']':
 294                 optional = False
 295                 var = var.strip()
 296                 var = var.strip('[')
 297                 var = var.strip(']')
 298                 chunks = var.split()
 299                 if len(chunks) > 1:
 300                     var = var.split()[0]
 301
 302                 # Environment vars the same as flag names without
 303                 # the initial -'s and in UPPERCASE.
 304                 env = var.strip('-').upper()
 305                 if env in os.environ:
 306                     if not is_flag_already_in_argv(var):
 307                         value = os.environ[env]
 308                         SAVED_MESSAGES.append(f'Initialized from environment: {var} = {value}')
 309                         from string_utils import to_bool
 310
 311                         if len(chunks) == 1 and to_bool(value):
 312                             sys.argv.append(var)
 313                         elif len(chunks) > 1:
 314                             sys.argv.append(var)
 315                             sys.argv.append(value)
 316                 var = ''
 317                 env = ''
 318
 319
 320 def _augment_sys_argv_from_loadfile():
 321     """Internal.  Augment with arguments persisted in a saved file."""
 322
 323     loadfile = None
 324     saw_other_args = False
 325     grab_next_arg = False
 326     for arg in sys.argv[1:]:
 327         if 'config_loadfile' in arg:
 328             pieces = arg.split('=')
 329             if len(pieces) > 1:
 330                 loadfile = pieces[1]
 331             else:
 332                 grab_next_arg = True
 333         elif grab_next_arg:
 334             loadfile = arg
 335         else:
 336             saw_other_args = True
 337
 338     if loadfile is not None:
 339         if not os.path.exists(loadfile):
 340             raise Exception(
 341                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 342             )
 343         if saw_other_args:
 344             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 345         else:
 346             msg = f'Reading commandline arguments from {loadfile}.'
 347         print(msg, file=sys.stderr)
 348         SAVED_MESSAGES.append(msg)
 349
 350         with open(loadfile, 'r') as rf:
 351             newargs = rf.readlines()
 352         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 353         sys.argv += newargs
 354
 355
 356 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 357     """Main program should call this early in main().  Note that the
 358     :code:`bootstrap.initialize` wrapper takes care of this automatically.
 359     This should only be called once per program invocation.
 360
 361     """
 362     global CONFIG_PARSE_CALLED
 363     if CONFIG_PARSE_CALLED:
 364         return config
 365
 366     # If we're about to do the usage message dump, put the main
 367     # module's argument group last in the list (if possible) so that
 368     # when the user passes -h or --help, it will be visible on the
 369     # screen w/o scrolling.
 370     for arg in sys.argv:
 371         if arg in ('--help', '-h'):
 372             if entry_module is not None:
 373                 entry_module = os.path.basename(entry_module)
 374             ARGS._action_groups = _reorder_arg_action_groups_before_help(entry_module)
 375
 376     # Examine the environment for variables that match known flags.
 377     # For a flag called --example_flag the corresponding environment
 378     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 379     # these into sys.argv to be parsed.
 380     _augment_sys_argv_from_environment_variables()
 381
 382     # Look for loadfile and read/parse it if present.  This also
 383     # works by jamming these values onto sys.argv.
 384     _augment_sys_argv_from_loadfile()
 385
 386     # Parse (possibly augmented, possibly completely overwritten)
 387     # commandline args with argparse normally and populate config.
 388     known, unknown = ARGS.parse_known_args()
 389     config.update(vars(known))
 390
 391     # Reconstruct the argv with unrecognized flags for the benefit of
 392     # future argument parsers.  For example, unittest_main in python
 393     # has some of its own flags.  If we didn't recognize it, maybe
 394     # someone else will.
 395     if len(unknown) > 0:
 396         if config['config_rejects_unrecognized_arguments']:
 397             raise Exception(
 398                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 399             )
 400         SAVED_MESSAGES.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 401     sys.argv = sys.argv[:1] + unknown
 402
 403     # Check for savefile and populate it if requested.
 404     savefile = config['config_savefile']
 405     if savefile and len(savefile) > 0:
 406         with open(savefile, 'w') as wf:
 407             wf.write("\n".join(ORIG_ARGV[1:]))
 408
 409     # Also dump the config on stderr if requested.
 410     if config['config_dump']:
 411         dump_config()
 412
 413     CONFIG_PARSE_CALLED = True
 414     return config
 415
 416
 417 def has_been_parsed() -> bool:
 418     """Returns True iff the global config has already been parsed"""
 419     return CONFIG_PARSE_CALLED
 420
 421
 422 def dump_config():
 423     """Print the current config to stdout."""
 424     print("Global Configuration:", file=sys.stderr)
 425     pprint.pprint(config, stream=sys.stderr)
 426     print()
 427
 428
 429 def late_logging():
 430     """Log messages saved earlier now that logging has been initialized."""
 431     logger = logging.getLogger(__name__)
 432     logger.debug('Original commandline was: %s', ORIG_ARGV)
 433     for _ in SAVED_MESSAGES:
 434         logger.debug(_)