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 _augment_sys_argv_from_environment_variables():
 258     """Internal.  Look at the system environment for variables that match
 259     arg names.  This is done via some munging such that:
 260
 261     :code:`--argument_to_match`
 262
 263     ...is matched by:
 264
 265     :code:`ARGUMENT_TO_MATCH`
 266
 267     This allows programmers to set args via shell environment variables
 268     in lieu of passing them on the cmdline.
 269
 270     """
 271
 272     usage_message = ARGS.format_usage()
 273     optional = False
 274     var = ''
 275     for x in usage_message.split():
 276         if x[0] == '[':
 277             optional = True
 278         if optional:
 279             var += f'{x} '
 280             if x[-1] == ']':
 281                 optional = False
 282                 var = var.strip()
 283                 var = var.strip('[')
 284                 var = var.strip(']')
 285                 chunks = var.split()
 286                 if len(chunks) > 1:
 287                     var = var.split()[0]
 288
 289                 # Environment vars the same as flag names without
 290                 # the initial -'s and in UPPERCASE.
 291                 env = var.strip('-').upper()
 292                 if env in os.environ:
 293                     if not is_flag_already_in_argv(var):
 294                         value = os.environ[env]
 295                         SAVED_MESSAGES.append(f'Initialized from environment: {var} = {value}')
 296                         from string_utils import to_bool
 297
 298                         if len(chunks) == 1 and to_bool(value):
 299                             sys.argv.append(var)
 300                         elif len(chunks) > 1:
 301                             sys.argv.append(var)
 302                             sys.argv.append(value)
 303                 var = ''
 304                 env = ''
 305
 306
 307 def _augment_sys_argv_from_loadfile():
 308     """Internal.  Augment with arguments persisted in a saved file."""
 309
 310     loadfile = None
 311     saw_other_args = False
 312     grab_next_arg = False
 313     for arg in sys.argv[1:]:
 314         if 'config_loadfile' in arg:
 315             pieces = arg.split('=')
 316             if len(pieces) > 1:
 317                 loadfile = pieces[1]
 318             else:
 319                 grab_next_arg = True
 320         elif grab_next_arg:
 321             loadfile = arg
 322         else:
 323             saw_other_args = True
 324
 325     if loadfile is not None:
 326         if not os.path.exists(loadfile):
 327             raise Exception(
 328                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 329             )
 330         if saw_other_args:
 331             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 332         else:
 333             msg = f'Reading commandline arguments from {loadfile}.'
 334         print(msg, file=sys.stderr)
 335         SAVED_MESSAGES.append(msg)
 336
 337         with open(loadfile, 'r') as rf:
 338             newargs = rf.readlines()
 339         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 340         sys.argv += newargs
 341
 342
 343 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 344     """Main program should call this early in main().  Note that the
 345     :code:`bootstrap.initialize` wrapper takes care of this automatically.
 346     This should only be called once per program invocation.
 347
 348     """
 349     global CONFIG_PARSE_CALLED
 350     if CONFIG_PARSE_CALLED:
 351         return config
 352
 353     # If we're about to do the usage message dump, put the main
 354     # module's argument group last in the list (if possible) so that
 355     # when the user passes -h or --help, it will be visible on the
 356     # screen w/o scrolling.
 357     for arg in sys.argv:
 358         if arg in ('--help', '-h'):
 359             if entry_module is not None:
 360                 entry_module = os.path.basename(entry_module)
 361             ARGS._action_groups = _reorder_arg_action_groups_before_help(entry_module)
 362
 363     # Examine the environment for variables that match known flags.
 364     # For a flag called --example_flag the corresponding environment
 365     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 366     # these into sys.argv to be parsed.
 367     _augment_sys_argv_from_environment_variables()
 368
 369     # Look for loadfile and read/parse it if present.  This also
 370     # works by jamming these values onto sys.argv.
 371     _augment_sys_argv_from_loadfile()
 372
 373     # Parse (possibly augmented, possibly completely overwritten)
 374     # commandline args with argparse normally and populate config.
 375     known, unknown = ARGS.parse_known_args()
 376     config.update(vars(known))
 377
 378     # Reconstruct the argv with unrecognized flags for the benefit of
 379     # future argument parsers.  For example, unittest_main in python
 380     # has some of its own flags.  If we didn't recognize it, maybe
 381     # someone else will.
 382     if len(unknown) > 0:
 383         if config['config_rejects_unrecognized_arguments']:
 384             raise Exception(
 385                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 386             )
 387         SAVED_MESSAGES.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 388     sys.argv = sys.argv[:1] + unknown
 389
 390     # Check for savefile and populate it if requested.
 391     savefile = config['config_savefile']
 392     if savefile and len(savefile) > 0:
 393         with open(savefile, 'w') as wf:
 394             wf.write("\n".join(ORIG_ARGV[1:]))
 395
 396     # Also dump the config on stderr if requested.
 397     if config['config_dump']:
 398         dump_config()
 399
 400     CONFIG_PARSE_CALLED = True
 401     return config
 402
 403
 404 def has_been_parsed() -> bool:
 405     """Returns True iff the global config has already been parsed"""
 406     return CONFIG_PARSE_CALLED
 407
 408
 409 def dump_config():
 410     """Print the current config to stdout."""
 411     print("Global Configuration:", file=sys.stderr)
 412     pprint.pprint(config, stream=sys.stderr)
 413     print()
 414
 415
 416 def late_logging():
 417     """Log messages saved earlier now that logging has been initialized."""
 418     logger = logging.getLogger(__name__)
 419     logger.debug('Original commandline was: %s', ORIG_ARGV)
 420     for _ in SAVED_MESSAGES:
 421         logger.debug(_)