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