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