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 saved configuration files, and zookeeper-based dynamic configurations.  This
   7 works across several modules.
   8
   9 Example usage:
  10
  11     In your file.py::
  12
  13         import config
  14
  15         parser = config.add_commandline_args(
  16             "Module",
  17             "Args related to module doing the thing.",
  18         )
  19         parser.add_argument(
  20             "--module_do_the_thing",
  21             type=bool,
  22             default=True,
  23             help="Should the module do the thing?"
  24         )
  25
  26     In your main.py::
  27
  28         import config
  29
  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
  41         def main() -> None:
  42             config.parse()   # Very important, this must be invoked!
  43
  44     If you set this up and remember to invoke config.parse(), all commandline
  45     arguments will play nicely together.  This is done automatically for you
  46     if you're using the :meth:`bootstrap.initialize` decorator on
  47     your program's entry point.  See :meth:`python_modules.bootstrap.initialize`
  48     for more details.::
  49
  50         import bootstrap
  51
  52         @bootstrap.initialize
  53         def main():
  54             whatever
  55
  56         if __name__ == '__main__':
  57             main()
  58
  59     Either way, you'll get this behavior from the commandline::
  60
  61         % main.py -h
  62         usage: main.py [-h]
  63                        [--module_do_the_thing MODULE_DO_THE_THING]
  64                        [--dry_run DRY_RUN]
  65
  66         Module:
  67           Args related to module doing the thing.
  68
  69           --module_do_the_thing MODULE_DO_THE_THING
  70                        Should the module do the thing?
  71
  72         Main:
  73           A program that does the thing
  74
  75           --dry_run
  76                        Should we really do the thing?
  77
  78     Arguments themselves should be accessed via
  79     :code:`config.config['arg_name']`.  e.g.::
  80
  81         if not config.config['dry_run']:
  82             module.do_the_thing()
  83 """
  84
  85 import argparse
  86 import logging
  87 import os
  88 import pprint
  89 import re
  90 import sys
  91 from typing import Any, Dict, List, Optional, Tuple
  92
  93 from kazoo.client import KazooClient
  94 from kazoo.protocol.states import WatchedEvent
  95
  96 import scott_secrets
  97
  98 # This module is commonly used by others in here and should avoid
  99 # taking any unnecessary dependencies back on them.
 100
 101 # Make a copy of the original program arguments immediately upon module load.
 102 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
 103 ORIG_ARGV: List[str] = sys.argv.copy()
 104
 105
 106 class OptionalRawFormatter(argparse.HelpFormatter):
 107     """This formatter has the same bahavior as the normal argparse text
 108     formatter except when the help text of an argument begins with
 109     "RAW|".  In that case, the line breaks are preserved and the text
 110     is not wrapped.
 111
 112     Use this, for example, when you need the helptext of an argument
 113     to have its spacing preserved exactly, e.g.::
 114
 115         args.add_argument(
 116             '--mode',
 117             type=str,
 118             default='PLAY',
 119             choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
 120             metavar='MODE',
 121             help='''RAW|Our mode of operation.  One of:
 122
 123                 PLAY = play wordle with me!  Pick a random solution or
 124                        specify a solution with --template.
 125
 126                CHEAT = given a --template and, optionally, --letters_in_word
 127                        and/or --letters_to_avoid, return the best guess word;
 128
 129             AUTOPLAY = given a complete word in --template, guess it step
 130                        by step showing work;
 131
 132             SELFTEST = autoplay every possible solution keeping track of
 133                        wins/losses and average number of guesses;
 134
 135           PRECOMPUTE = populate hash table with optimal guesses.
 136             ''',
 137         )
 138     """
 139
 140     def _split_lines(self, text, width):
 141         if text.startswith('RAW|'):
 142             return text[4:].splitlines()
 143         return argparse.HelpFormatter._split_lines(self, text, width)
 144
 145
 146 # A global argparser that we will collect arguments in.  Each module (including
 147 # us) will add arguments to a separate argument group.
 148 ARGS = argparse.ArgumentParser(
 149     description=None,
 150     formatter_class=OptionalRawFormatter,
 151     fromfile_prefix_chars="@",
 152     epilog=f'{PROGRAM_NAME} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
 153     # I don't fully understand why but when loaded by sphinx sometimes
 154     # the same module is loaded many times causing any arguments it
 155     # registers via module-level code to be redefined.  Work around
 156     # this iff the program is 'sphinx-build'
 157     conflict_handler='resolve' if PROGRAM_NAME == 'sphinx-build' else 'error',
 158 )
 159
 160 # Arguments specific to config.py.  Other users should get their own group by
 161 # invoking config.add_commandline_args.
 162 GROUP = ARGS.add_argument_group(
 163     f'Global Config ({__file__})',
 164     'Args that control the global config itself; how meta!',
 165 )
 166 GROUP.add_argument(
 167     '--config_loadfile',
 168     metavar='FILENAME',
 169     default=None,
 170     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to those passed via the commandline.  Note that if the given path begins with "zk:" then it is interpreted as a zookeeper path instead of as a filesystem path.  When loading config from zookeeper, any argument with the string "dynamic" in the name (e.g. --module_dynamic_url) may be modified at runtime by changes made to zookeeper (using --config_savefile=zk:path).  You should therefore either write your code to handle dynamic argument changes or avoid naming arguments "dynamic" if you use zookeeper configuration paths.',
 171 )
 172 GROUP.add_argument(
 173     '--config_dump',
 174     default=False,
 175     action='store_true',
 176     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup time.',
 177 )
 178 GROUP.add_argument(
 179     '--config_savefile',
 180     type=str,
 181     metavar='FILENAME',
 182     default=None,
 183     help='Populate a config file (compatible with --config_loadfile) with the given path for later use.  If the given path begins with "zk:" it is interpreted as a zookeeper path instead of a filesystem path.  When updating zookeeper-based configs, all running programs that read their configuration from zookeeper (via --config_loadfile=zk:path) at startup time will see their configuration dynamically updated; flags with "dynamic" in their names (e.g. --my_dynamic_flag) may have their values changed.  You should therefore either write your code to handle dynamic argument changes or avoid naming arguments "dynamic" if you use zookeeper configuration paths.',
 184 )
 185 GROUP.add_argument(
 186     '--config_rejects_unrecognized_arguments',
 187     default=False,
 188     action='store_true',
 189     help='If present, config will raise an exception if it doesn\'t recognize an argument.  The default behavior is to ignore unknown arguments so as to allow interoperability with programs that want to use their own argparse calls to parse their own, separate commandline args.',
 190 )
 191
 192
 193 class Config:
 194     """
 195     Everything in the config module used to be module-level functions and
 196     variables but it made the code ugly and harder to maintain.  Now, this
 197     class does the heavy lifting.  We still rely on some globals, though:
 198
 199         ARGS and GROUP to interface with argparse
 200         PROGRAM_NAME stores argv[0] close to program invocation
 201         ORIG_ARGV stores the original argv list close to program invocation
 202         CONFIG and config: hold the (singleton) instance of this class.
 203
 204     """
 205
 206     def __init__(self):
 207         # Has our parse() method been invoked yet?
 208         self.config_parse_called = False
 209
 210         # A configuration dictionary that will contain parsed
 211         # arguments.  This is the data that is most interesting to our
 212         # callers as it will hold the configuration result.
 213         self.config: Dict[str, Any] = {}
 214
 215         # Defer logging messages until later when logging has been
 216         # initialized.
 217         self.saved_messages: List[str] = []
 218
 219         # A zookeeper client that is lazily created so as to not incur
 220         # the latency of connecting to zookeeper for programs that are
 221         # not reading or writing their config data into zookeeper.
 222         self.zk: Optional[KazooClient] = None
 223
 224         # Per known zk file, what is the max version we have seen?
 225         self.max_version: Dict[str, int] = {}
 226
 227     def __getitem__(self, key: str) -> Optional[Any]:
 228         """If someone uses []'s on us, pass it onto self.config."""
 229         return self.config.get(key, None)
 230
 231     def __setitem__(self, key: str, value: Any) -> None:
 232         self.config[key] = value
 233
 234     def __contains__(self, key: str) -> bool:
 235         return key in self.config
 236
 237     @staticmethod
 238     def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
 239         """Create a new context for arguments and return a handle.
 240
 241         Args:
 242             title: A title for your module's commandline arguments group.
 243             description: A helpful description of your module.
 244
 245         Returns:
 246             An argparse._ArgumentGroup to be populated by the caller.
 247         """
 248         return ARGS.add_argument_group(title, description)
 249
 250     @staticmethod
 251     def overwrite_argparse_epilog(msg: str) -> None:
 252         """Allows your code to override the default epilog created by
 253         argparse.
 254
 255         Args:
 256             msg: The epilog message to substitute for the default.
 257         """
 258         ARGS.epilog = msg
 259
 260     @staticmethod
 261     def is_flag_already_in_argv(var: str) -> bool:
 262         """Returns true if a particular flag is passed on the commandline
 263         and false otherwise.
 264
 265         Args:
 266             var: The flag to search for.
 267         """
 268         for _ in sys.argv:
 269             if var in _:
 270                 return True
 271         return False
 272
 273     @staticmethod
 274     def print_usage() -> None:
 275         """Prints the normal help usage message out."""
 276         ARGS.print_help()
 277
 278     @staticmethod
 279     def usage() -> str:
 280         """
 281         Returns:
 282             program usage help text as a string.
 283         """
 284         return ARGS.format_usage()
 285
 286     @staticmethod
 287     def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
 288         """Internal.  Used to reorder the arguments before dumping out a
 289         generated help string such that the main program's arguments come
 290         last.
 291
 292         """
 293         reordered_action_groups = []
 294         for grp in ARGS._action_groups:
 295             if entry_module is not None and entry_module in grp.title:  # type: ignore
 296                 reordered_action_groups.append(grp)
 297             elif PROGRAM_NAME in GROUP.title:  # type: ignore
 298                 reordered_action_groups.append(grp)
 299             else:
 300                 reordered_action_groups.insert(0, grp)
 301         return reordered_action_groups
 302
 303     @staticmethod
 304     def _parse_arg_into_env(arg: str) -> Optional[Tuple[str, str, List[str]]]:
 305         """Internal helper to parse commandling args into environment vars."""
 306
 307         arg = arg.strip()
 308         if not arg.startswith('['):
 309             return None
 310         arg = arg.strip('[')
 311         if not arg.endswith(']'):
 312             return None
 313         arg = arg.strip(']')
 314
 315         chunks = arg.split()
 316         if len(chunks) > 1:
 317             var = arg.split()[0]
 318         else:
 319             var = arg
 320
 321         # Environment vars the same as flag names without
 322         # the initial -'s and in UPPERCASE.
 323         env = var.upper()
 324         while env[0] == '-':
 325             env = env[1:]
 326         return arg, env, chunks
 327
 328     def _augment_sys_argv_from_environment_variables(self):
 329         """Internal.  Look at the system environment for variables that match
 330         commandline arg names.  This is done via some munging such that:
 331
 332         :code:`--argument_to_match`
 333
 334         ...is matched by:
 335
 336         :code:`ARGUMENT_TO_MATCH`
 337
 338         This allows users to set args via shell environment variables
 339         in lieu of passing them on the cmdline.
 340
 341         """
 342         usage_message = Config.usage()
 343         optional = False
 344         arg = ''
 345
 346         # Foreach valid optional commandline option (chunk) generate
 347         # its analogous environment variable.
 348         for chunk in usage_message.split():
 349             if chunk[0] == '[':
 350                 optional = True
 351             if optional:
 352                 arg += f'{chunk} '
 353                 if chunk[-1] == ']':
 354                     optional = False
 355                     _ = Config._parse_arg_into_env(arg)
 356                     if _:
 357                         var, env, chunks = _
 358                         if env in os.environ:
 359                             if not Config.is_flag_already_in_argv(var):
 360                                 value = os.environ[env]
 361                                 self.saved_messages.append(
 362                                     f'Initialized from environment: {var} = {value}'
 363                                 )
 364                                 from string_utils import to_bool
 365
 366                                 if len(chunks) == 1 and to_bool(value):
 367                                     sys.argv.append(arg)
 368                                 elif len(chunks) > 1:
 369                                     sys.argv.append(arg)
 370                                     sys.argv.append(value)
 371                     arg = ''
 372
 373     def _process_dynamic_args(self, event: WatchedEvent):
 374         assert self.zk
 375         logger = logging.getLogger(__name__)
 376         contents, meta = self.zk.get(event.path, watch=self._process_dynamic_args)
 377         logger.debug('Update for %s at version=%d.', event.path, meta.version)
 378         logger.debug(
 379             'Max known version for %s is %d.', event.path, self.max_version.get(event.path, 0)
 380         )
 381         if meta.version > self.max_version.get(event.path, 0):
 382             self.max_version[event.path] = meta.version
 383             contents = contents.decode()
 384             temp_argv = []
 385             for arg in contents.split():
 386                 if 'dynamic' in arg:
 387                     temp_argv.append(arg)
 388                     logger.info("Updating %s from zookeeper async config change.", arg)
 389             if len(temp_argv) > 0:
 390                 old_argv = sys.argv
 391                 sys.argv = temp_argv
 392                 known, _ = ARGS.parse_known_args()
 393                 sys.argv = old_argv
 394                 self.config.update(vars(known))
 395
 396     def _augment_sys_argv_from_loadfile(self):
 397         """Internal.  Augment with arguments persisted in a saved file."""
 398
 399         loadfile = None
 400         saw_other_args = False
 401         grab_next_arg = False
 402         for arg in sys.argv[1:]:
 403             if 'config_loadfile' in arg:
 404                 pieces = arg.split('=')
 405                 if len(pieces) > 1:
 406                     loadfile = pieces[1]
 407                 else:
 408                     grab_next_arg = True
 409             elif grab_next_arg:
 410                 loadfile = arg
 411             else:
 412                 saw_other_args = True
 413
 414         if loadfile is not None:
 415             zkpath = None
 416             if loadfile[:3] == 'zk:':
 417                 try:
 418                     if self.zk is None:
 419                         self.zk = KazooClient(
 420                             hosts=scott_secrets.ZOOKEEPER_NODES,
 421                             use_ssl=True,
 422                             verify_certs=False,
 423                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 424                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
 425                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 426                         )
 427                         self.zk.start()
 428                     zkpath = loadfile[3:]
 429                     if not zkpath.startswith('/config/'):
 430                         zkpath = '/config/' + zkpath
 431                         zkpath = re.sub(r'//+', '/', zkpath)
 432                     if not self.zk.exists(zkpath):
 433                         raise Exception(
 434                             f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
 435                         )
 436                 except Exception as e:
 437                     raise Exception(
 438                         f'ERROR: Error talking with zookeeper while looking for {loadfile}'
 439                     ) from e
 440             elif not os.path.exists(loadfile):
 441                 raise Exception(
 442                     f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 443                 )
 444
 445             if saw_other_args:
 446                 msg = f'Augmenting commandline arguments with those from {loadfile}.'
 447             else:
 448                 msg = f'Reading commandline arguments from {loadfile}.'
 449             print(msg, file=sys.stderr)
 450             self.saved_messages.append(msg)
 451
 452             newargs = []
 453             if zkpath:
 454                 try:
 455                     assert self.zk
 456                     contents, meta = self.zk.get(zkpath, watch=self._process_dynamic_args)
 457                     contents = contents.decode()
 458                     newargs = [
 459                         arg.strip('\n')
 460                         for arg in contents.split('\n')
 461                         if 'config_savefile' not in arg
 462                     ]
 463                     self.saved_messages.append(f'Setting {zkpath}\'s max_version to {meta.version}')
 464                     self.max_version[zkpath] = meta.version
 465                 except Exception as e:
 466                     raise Exception(f'Error reading {zkpath} from zookeeper.') from e
 467                 self.saved_messages.append(f'Loaded config from zookeeper from {zkpath}')
 468             else:
 469                 with open(loadfile, 'r') as rf:
 470                     newargs = rf.readlines()
 471                 newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 472             sys.argv += newargs
 473
 474     def dump_config(self):
 475         """Print the current config to stdout."""
 476         print("Global Configuration:", file=sys.stderr)
 477         pprint.pprint(self.config, stream=sys.stderr)
 478         print()
 479
 480     def parse(self, entry_module: Optional[str]) -> Dict[str, Any]:
 481         """Main program should call this early in main().  Note that the
 482         :code:`bootstrap.initialize` wrapper takes care of this automatically.
 483         This should only be called once per program invocation.
 484
 485         """
 486         if self.config_parse_called:
 487             return self.config
 488
 489         # If we're about to do the usage message dump, put the main
 490         # module's argument group last in the list (if possible) so that
 491         # when the user passes -h or --help, it will be visible on the
 492         # screen w/o scrolling.
 493         for arg in sys.argv:
 494             if arg in ('--help', '-h'):
 495                 if entry_module is not None:
 496                     entry_module = os.path.basename(entry_module)
 497                 ARGS._action_groups = Config._reorder_arg_action_groups_before_help(entry_module)
 498
 499         # Examine the environment for variables that match known flags.
 500         # For a flag called --example_flag the corresponding environment
 501         # variable would be called EXAMPLE_FLAG.  If found, hackily add
 502         # these into sys.argv to be parsed.
 503         self._augment_sys_argv_from_environment_variables()
 504
 505         # Look for loadfile and read/parse it if present.  This also
 506         # works by jamming these values onto sys.argv.
 507         self._augment_sys_argv_from_loadfile()
 508
 509         # Parse (possibly augmented, possibly completely overwritten)
 510         # commandline args with argparse normally and populate config.
 511         known, unknown = ARGS.parse_known_args()
 512         self.config.update(vars(known))
 513
 514         # Reconstruct the argv with unrecognized flags for the benefit of
 515         # future argument parsers.  For example, unittest_main in python
 516         # has some of its own flags.  If we didn't recognize it, maybe
 517         # someone else will.
 518         if len(unknown) > 0:
 519             if config['config_rejects_unrecognized_arguments']:
 520                 raise Exception(
 521                     f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 522                 )
 523             self.saved_messages.append(
 524                 f'Config encountered unrecognized commandline arguments: {unknown}'
 525             )
 526         sys.argv = sys.argv[:1] + unknown
 527
 528         # Check for savefile and populate it if requested.
 529         savefile = config['config_savefile']
 530         if savefile and len(savefile) > 0:
 531             data = '\n'.join(ORIG_ARGV[1:])
 532             if savefile[:3] == 'zk:':
 533                 zkpath = savefile[3:]
 534                 if not zkpath.startswith('/config/'):
 535                     zkpath = '/config/' + zkpath
 536                     zkpath = re.sub(r'//+', '/', zkpath)
 537                 try:
 538                     if not self.zk:
 539                         self.zk = KazooClient(
 540                             hosts=scott_secrets.ZOOKEEPER_NODES,
 541                             use_ssl=True,
 542                             verify_certs=False,
 543                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 544                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
 545                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 546                         )
 547                         self.zk.start()
 548                     data = data.encode()
 549                     if len(data) > 1024 * 1024:
 550                         raise Exception(f'Saved args are too large!  ({len(data)} bytes)')
 551                     if not self.zk.exists(zkpath):
 552                         self.zk.create(zkpath, data)
 553                         self.saved_messages.append(
 554                             f'Just created {zkpath}; setting its max_version to 0'
 555                         )
 556                         self.max_version[zkpath] = 0
 557                     else:
 558                         meta = self.zk.set(zkpath, data)
 559                         self.saved_messages.append(
 560                             f'Setting {zkpath}\'s max_version to {meta.version}'
 561                         )
 562                         self.max_version[zkpath] = meta.version
 563                 except Exception as e:
 564                     raise Exception(f'Failed to create zookeeper path {zkpath}') from e
 565                 self.saved_messages.append(f'Saved config to zookeeper in {zkpath}')
 566             else:
 567                 with open(savefile, 'w') as wf:
 568                     wf.write(data)
 569
 570         # Also dump the config on stderr if requested.
 571         if config['config_dump']:
 572             self.dump_config()
 573
 574         self.config_parse_called = True
 575         return self.config
 576
 577     def has_been_parsed(self) -> bool:
 578         """Returns True iff the global config has already been parsed"""
 579         return self.config_parse_called
 580
 581     def late_logging(self):
 582         """Log messages saved earlier now that logging has been initialized."""
 583         logger = logging.getLogger(__name__)
 584         logger.debug('Original commandline was: %s', ORIG_ARGV)
 585         for _ in self.saved_messages:
 586             logger.debug(_)
 587
 588
 589 # A global singleton instance of the Config class.
 590 CONFIG = Config()
 591
 592 # A lot of client code uses config.config['whatever'] to lookup
 593 # configuration so to preserve this we make this, config.config, with
 594 # a __getitem__ method on it.
 595 config = CONFIG
 596
 597 # Config didn't use to be a class; it was a mess of module-level
 598 # functions and data.  The functions below preserve the old interface
 599 # so that existing clients do not need to be changed.  As you can see,
 600 # they mostly just thunk into the config class.
 601
 602
 603 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
 604     """Create a new context for arguments and return a handle.  An alias
 605     for config.config.add_commandline_args.
 606
 607     Args:
 608         title: A title for your module's commandline arguments group.
 609         description: A helpful description of your module.
 610
 611     Returns:
 612         An argparse._ArgumentGroup to be populated by the caller.
 613     """
 614     return CONFIG.add_commandline_args(title, description)
 615
 616
 617 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 618     """Main program should call this early in main().  Note that the
 619     :code:`bootstrap.initialize` wrapper takes care of this automatically.
 620     This should only be called once per program invocation.  Subsequent
 621     calls do not reparse the configuration settings but rather just
 622     return the current state.
 623     """
 624     return CONFIG.parse(entry_module)
 625
 626
 627 def has_been_parsed() -> bool:
 628     """Returns True iff the global config has already been parsed"""
 629     return CONFIG.has_been_parsed()
 630
 631
 632 def late_logging() -> None:
 633     """Log messages saved earlier now that logging has been initialized."""
 634     CONFIG.late_logging()
 635
 636
 637 def dump_config() -> None:
 638     """Print the current config to stdout."""
 639     CONFIG.dump_config()
 640
 641
 642 def overwrite_argparse_epilog(msg: str) -> None:
 643     """Allows your code to override the default epilog created by
 644     argparse.
 645
 646     Args:
 647         msg: The epilog message to substitute for the default.
 648     """
 649     Config.overwrite_argparse_epilog(msg)
 650
 651
 652 def is_flag_already_in_argv(var: str) -> bool:
 653     """Returns true if a particular flag is passed on the commandline
 654     and false otherwise.
 655
 656     Args:
 657         var: The flag to search for.
 658     """
 659     return Config.is_flag_already_in_argv(var)
 660
 661
 662 def print_usage() -> None:
 663     """Prints the normal help usage message out."""
 664     Config.print_usage()
 665
 666
 667 def usage() -> str:
 668     """
 669     Returns:
 670         program usage help text as a string.
 671     """
 672     return Config.usage()