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