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 import argparse
  84 import logging
  85 import os
  86 import pprint
  87 import re
  88 import sys
  89 from typing import Any, Dict, List, Optional
  90
  91 from kazoo.client import KazooClient
  92
  93 import scott_secrets
  94
  95 # This module is commonly used by others in here and should avoid
  96 # taking any unnecessary dependencies back on them.
  97
  98 # Make a copy of the original program arguments.
  99 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
 100 ORIG_ARGV: List[str] = sys.argv.copy()
 101
 102
 103 class OptionalRawFormatter(argparse.HelpFormatter):
 104     """This formatter has the same bahavior as the normal argparse text
 105     formatter except when the help text of an argument begins with
 106     "RAW|".  In that case, the line breaks are preserved and the text
 107     is not wrapped.
 108
 109     Use this, for example, when you need the helptext of an argument
 110     to have its spacing preserved exactly, e.g.::
 111
 112         args.add_argument(
 113             '--mode',
 114             type=str,
 115             default='PLAY',
 116             choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
 117             metavar='MODE',
 118             help='''RAW|Our mode of operation.  One of:
 119
 120                 PLAY = play wordle with me!  Pick a random solution or
 121                        specify a solution with --template.
 122
 123                CHEAT = given a --template and, optionally, --letters_in_word
 124                        and/or --letters_to_avoid, return the best guess word;
 125
 126             AUTOPLAY = given a complete word in --template, guess it step
 127                        by step showing work;
 128
 129             SELFTEST = autoplay every possible solution keeping track of
 130                        wins/losses and average number of guesses;
 131
 132           PRECOMPUTE = populate hash table with optimal guesses.
 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 GROUP = ARGS.add_argument_group(
 156     f'Global Config ({__file__})',
 157     'Args that control the global config itself; how meta!',
 158 )
 159 GROUP.add_argument(
 160     '--config_loadfile',
 161     metavar='FILENAME',
 162     default=None,
 163     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.  Note that if this begins with zk: the path is interpreted as a zookeeper path instead of a filesystem path.',
 164 )
 165 GROUP.add_argument(
 166     '--config_dump',
 167     default=False,
 168     action='store_true',
 169     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 170 )
 171 GROUP.add_argument(
 172     '--config_savefile',
 173     type=str,
 174     metavar='FILENAME',
 175     default=None,
 176     help='Populate config file compatible with --config_loadfile to save global config for later use.  Note that if this begins with zk: the path is interpreted as a zookeeper oath instead of a filesystem path.',
 177 )
 178 GROUP.add_argument(
 179     '--config_rejects_unrecognized_arguments',
 180     default=False,
 181     action='store_true',
 182     help='If present, config will raise an exception if it doesn\'t recognize an argument.  The default behavior is to ignore this so as to allow interoperability with programs that want to use their own argparse calls to parse their own, separate commandline args.',
 183 )
 184
 185
 186 class Config:
 187     """
 188     Everything in the config module used to be module-level functions and
 189     variables but it made the code ugly and harder to maintain.  Now, this
 190     class does the heavy lifting.  We still rely on some globals, though:
 191
 192         ARGS and GROUP to interface with argparse
 193         PROGRAM_NAME stores argv[0] close to program invocation
 194         ORIG_ARGV stores the original argv list close to program invocation
 195         CONFIG and config: hold the (singleton) instance of this class.
 196
 197     """
 198
 199     def __init__(self):
 200         # Has our parse() method been invoked yet?
 201         self.config_parse_called = False
 202
 203         # A configuration dictionary that will contain parsed
 204         # arguments.  This is the data that is most interesting to our
 205         # callers as it will hold the configuration result.
 206         self.config: Dict[str, Any] = {}
 207
 208         # Defer logging messages until later when logging has been
 209         # initialized.
 210         self.saved_messages: List[str] = []
 211
 212         # A zookeeper client that is lazily created so as to not incur
 213         # the latency of connecting to zookeeper for programs that are
 214         # not reading or writing their config data into zookeeper.
 215         self.zk: Optional[KazooClient] = None
 216
 217     def __getitem__(self, key: str) -> Optional[Any]:
 218         """If someone uses []'s on us, pass it onto self.config."""
 219         return self.config.get(key, None)
 220
 221     def __setitem__(self, key: str, value: Any) -> None:
 222         self.config[key] = value
 223
 224     def __contains__(self, key: str) -> bool:
 225         return key in self.config
 226
 227     @staticmethod
 228     def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
 229         """Create a new context for arguments and return a handle.
 230
 231         Args:
 232             title: A title for your module's commandline arguments group.
 233             description: A helpful description of your module.
 234
 235         Returns:
 236             An argparse._ArgumentGroup to be populated by the caller.
 237         """
 238         return ARGS.add_argument_group(title, description)
 239
 240     @staticmethod
 241     def overwrite_argparse_epilog(msg: str) -> None:
 242         """Allows your code to override the default epilog created by
 243         argparse.
 244
 245         Args:
 246             msg: The epilog message to substitute for the default.
 247         """
 248         ARGS.epilog = msg
 249
 250     @staticmethod
 251     def is_flag_already_in_argv(var: str) -> bool:
 252         """Returns true if a particular flag is passed on the commandline
 253         and false otherwise.
 254
 255         Args:
 256             var: The flag to search for.
 257         """
 258         for _ in sys.argv:
 259             if var in _:
 260                 return True
 261         return False
 262
 263     @staticmethod
 264     def print_usage() -> None:
 265         """Prints the normal help usage message out."""
 266         ARGS.print_help()
 267
 268     @staticmethod
 269     def usage() -> str:
 270         """
 271         Returns:
 272             program usage help text as a string.
 273         """
 274         return ARGS.format_usage()
 275
 276     @staticmethod
 277     def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
 278         """Internal.  Used to reorder the arguments before dumping out a
 279         generated help string such that the main program's arguments come
 280         last.
 281
 282         """
 283         reordered_action_groups = []
 284         for grp in ARGS._action_groups:
 285             if entry_module is not None and entry_module in grp.title:  # type: ignore
 286                 reordered_action_groups.append(grp)
 287             elif PROGRAM_NAME in GROUP.title:  # type: ignore
 288                 reordered_action_groups.append(grp)
 289             else:
 290                 reordered_action_groups.insert(0, grp)
 291         return reordered_action_groups
 292
 293     def _augment_sys_argv_from_environment_variables(self):
 294         """Internal.  Look at the system environment for variables that match
 295         arg names.  This is done via some munging such that:
 296
 297         :code:`--argument_to_match`
 298
 299         ...is matched by:
 300
 301         :code:`ARGUMENT_TO_MATCH`
 302
 303         This allows programmers to set args via shell environment variables
 304         in lieu of passing them on the cmdline.
 305
 306         """
 307         usage_message = Config.usage()
 308         optional = False
 309         var = ''
 310
 311         # Foreach valid optional commandline option (line) generate
 312         # its analogous environment variable.
 313         for line in usage_message.split():
 314             if line[0] == '[':
 315                 optional = True
 316             if optional:
 317                 var += f'{line} '
 318                 if line[-1] == ']':
 319                     optional = False
 320                     var = var.strip()
 321                     var = var.strip('[')
 322                     var = var.strip(']')
 323                     chunks = var.split()
 324                     if len(chunks) > 1:
 325                         var = var.split()[0]
 326
 327                     # Environment vars the same as flag names without
 328                     # the initial -'s and in UPPERCASE.
 329                     env = var.upper()
 330                     while env[0] == '-':
 331                         env = env[1:]
 332
 333                     # Do we see that environment varaible?
 334                     if env in os.environ:
 335                         if not Config.is_flag_already_in_argv(var):
 336                             value = os.environ[env]
 337                             self.saved_messages.append(
 338                                 f'Initialized from environment: {var} = {value}'
 339                             )
 340                             from string_utils import to_bool
 341
 342                             if len(chunks) == 1 and to_bool(value):
 343                                 sys.argv.append(var)
 344                             elif len(chunks) > 1:
 345                                 sys.argv.append(var)
 346                                 sys.argv.append(value)
 347                     var = ''
 348                     env = ''
 349                     chunks = []
 350
 351     def _augment_sys_argv_from_loadfile(self):
 352         """Internal.  Augment with arguments persisted in a saved file."""
 353
 354         loadfile = None
 355         saw_other_args = False
 356         grab_next_arg = False
 357         for arg in sys.argv[1:]:
 358             if 'config_loadfile' in arg:
 359                 pieces = arg.split('=')
 360                 if len(pieces) > 1:
 361                     loadfile = pieces[1]
 362                 else:
 363                     grab_next_arg = True
 364             elif grab_next_arg:
 365                 loadfile = arg
 366             else:
 367                 saw_other_args = True
 368
 369         if loadfile is not None:
 370             zkpath = None
 371             if loadfile[:3] == 'zk:':
 372                 try:
 373                     if self.zk is None:
 374                         self.zk = KazooClient(
 375                             hosts=scott_secrets.ZOOKEEPER_NODES,
 376                             use_ssl=True,
 377                             verify_certs=False,
 378                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 379                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
 380                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 381                         )
 382                         self.zk.start()
 383                     zkpath = loadfile[3:]
 384                     if not zkpath.startswith('/config/'):
 385                         zkpath = '/config/' + zkpath
 386                         zkpath = re.sub(r'//+', '/', zkpath)
 387                     if not self.zk.exists(zkpath):
 388                         raise Exception(
 389                             f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
 390                         )
 391                 except Exception as e:
 392                     raise Exception(
 393                         f'ERROR: Error talking with zookeeper while looking for {loadfile}'
 394                     ) from e
 395             elif not os.path.exists(loadfile):
 396                 raise Exception(
 397                     f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 398                 )
 399
 400             if saw_other_args:
 401                 msg = f'Augmenting commandline arguments with those from {loadfile}.'
 402             else:
 403                 msg = f'Reading commandline arguments from {loadfile}.'
 404             print(msg, file=sys.stderr)
 405             self.saved_messages.append(msg)
 406
 407             newargs = []
 408             if zkpath:
 409                 try:
 410                     assert self.zk
 411                     contents = self.zk.get(zkpath)[0]
 412                     contents = contents.decode()
 413                     newargs = [
 414                         arg.strip('\n')
 415                         for arg in contents.split('\n')
 416                         if 'config_savefile' not in arg
 417                     ]
 418                     size = sys.getsizeof(newargs)
 419                     if size > 1024 * 1024:
 420                         raise Exception(f'Saved args are too large!  ({size} bytes)')
 421                 except Exception as e:
 422                     raise Exception(f'Error reading {zkpath} from zookeeper.') from e
 423                 self.saved_messages.append(f'Loaded config from zookeeper from {zkpath}')
 424             else:
 425                 with open(loadfile, 'r') as rf:
 426                     newargs = rf.readlines()
 427                 newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 428             sys.argv += newargs
 429
 430     def dump_config(self):
 431         """Print the current config to stdout."""
 432         print("Global Configuration:", file=sys.stderr)
 433         pprint.pprint(self.config, stream=sys.stderr)
 434         print()
 435
 436     def parse(self, entry_module: Optional[str]) -> Dict[str, Any]:
 437         """Main program should call this early in main().  Note that the
 438         :code:`bootstrap.initialize` wrapper takes care of this automatically.
 439         This should only be called once per program invocation.
 440
 441         """
 442         if self.config_parse_called:
 443             return self.config
 444
 445         # If we're about to do the usage message dump, put the main
 446         # module's argument group last in the list (if possible) so that
 447         # when the user passes -h or --help, it will be visible on the
 448         # screen w/o scrolling.
 449         for arg in sys.argv:
 450             if arg in ('--help', '-h'):
 451                 if entry_module is not None:
 452                     entry_module = os.path.basename(entry_module)
 453                 ARGS._action_groups = Config._reorder_arg_action_groups_before_help(entry_module)
 454
 455         # Examine the environment for variables that match known flags.
 456         # For a flag called --example_flag the corresponding environment
 457         # variable would be called EXAMPLE_FLAG.  If found, hackily add
 458         # these into sys.argv to be parsed.
 459         self._augment_sys_argv_from_environment_variables()
 460
 461         # Look for loadfile and read/parse it if present.  This also
 462         # works by jamming these values onto sys.argv.
 463         self._augment_sys_argv_from_loadfile()
 464
 465         # Parse (possibly augmented, possibly completely overwritten)
 466         # commandline args with argparse normally and populate config.
 467         known, unknown = ARGS.parse_known_args()
 468         self.config.update(vars(known))
 469
 470         # Reconstruct the argv with unrecognized flags for the benefit of
 471         # future argument parsers.  For example, unittest_main in python
 472         # has some of its own flags.  If we didn't recognize it, maybe
 473         # someone else will.
 474         if len(unknown) > 0:
 475             if config['config_rejects_unrecognized_arguments']:
 476                 raise Exception(
 477                     f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 478                 )
 479             self.saved_messages.append(
 480                 f'Config encountered unrecognized commandline arguments: {unknown}'
 481             )
 482         sys.argv = sys.argv[:1] + unknown
 483
 484         # Check for savefile and populate it if requested.
 485         savefile = config['config_savefile']
 486         if savefile and len(savefile) > 0:
 487             data = '\n'.join(ORIG_ARGV[1:])
 488             if savefile[:3] == 'zk:':
 489                 zkpath = savefile[3:]
 490                 if not zkpath.startswith('/config/'):
 491                     zkpath = '/config/' + zkpath
 492                     zkpath = re.sub(r'//+', '/', zkpath)
 493                 try:
 494                     if not self.zk:
 495                         self.zk = KazooClient(
 496                             hosts=scott_secrets.ZOOKEEPER_NODES,
 497                             use_ssl=True,
 498                             verify_certs=False,
 499                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 500                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
 501                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
 502                         )
 503                         self.zk.start()
 504                     if not self.zk.exists(zkpath):
 505                         self.zk.create(zkpath, data.encode())
 506                     else:
 507                         self.zk.set(zkpath, data.encode())
 508                 except Exception as e:
 509                     raise Exception(f'Failed to create zookeeper path {zkpath}') from e
 510                 self.saved_messages.append(f'Saved config to zookeeper in {zkpath}')
 511             else:
 512                 with open(savefile, 'w') as wf:
 513                     wf.write(data)
 514
 515         # Also dump the config on stderr if requested.
 516         if config['config_dump']:
 517             self.dump_config()
 518
 519         self.config_parse_called = True
 520         return self.config
 521
 522     def has_been_parsed(self) -> bool:
 523         """Returns True iff the global config has already been parsed"""
 524         return self.config_parse_called
 525
 526     def late_logging(self):
 527         """Log messages saved earlier now that logging has been initialized."""
 528         logger = logging.getLogger(__name__)
 529         logger.debug('Original commandline was: %s', ORIG_ARGV)
 530         for _ in self.saved_messages:
 531             logger.debug(_)
 532
 533
 534 # A global singleton instance of the Config class.
 535 CONFIG = Config()
 536
 537 # A lot of client code uses config.config['whatever'] to lookup
 538 # configuration so to preserve this we make this, config.config, with
 539 # a __getitem__ method on it.
 540 config = CONFIG
 541
 542 # Config didn't use to be a class; it was a mess of module-level
 543 # functions and data.  The functions below preserve the old interface
 544 # so that existing clients do not need to be changed.  As you can see,
 545 # they mostly just thunk into the config class.
 546
 547
 548 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
 549     """Create a new context for arguments and return a handle.  An alias
 550     for config.config.add_commandline_args.
 551
 552     Args:
 553         title: A title for your module's commandline arguments group.
 554         description: A helpful description of your module.
 555
 556     Returns:
 557         An argparse._ArgumentGroup to be populated by the caller.
 558     """
 559     return CONFIG.add_commandline_args(title, description)
 560
 561
 562 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 563     """Main program should call this early in main().  Note that the
 564     :code:`bootstrap.initialize` wrapper takes care of this automatically.
 565     This should only be called once per program invocation.  Subsequent
 566     calls do not reparse the configuration settings but rather just
 567     return the current state.
 568     """
 569     return CONFIG.parse(entry_module)
 570
 571
 572 def has_been_parsed() -> bool:
 573     """Returns True iff the global config has already been parsed"""
 574     return CONFIG.has_been_parsed()
 575
 576
 577 def late_logging() -> None:
 578     """Log messages saved earlier now that logging has been initialized."""
 579     CONFIG.late_logging()
 580
 581
 582 def dump_config() -> None:
 583     """Print the current config to stdout."""
 584     CONFIG.dump_config()
 585
 586
 587 def overwrite_argparse_epilog(msg: str) -> None:
 588     """Allows your code to override the default epilog created by
 589     argparse.
 590
 591     Args:
 592         msg: The epilog message to substitute for the default.
 593     """
 594     Config.overwrite_argparse_epilog(msg)
 595
 596
 597 def is_flag_already_in_argv(var: str) -> bool:
 598     """Returns true if a particular flag is passed on the commandline
 599     and false otherwise.
 600
 601     Args:
 602         var: The flag to search for.
 603     """
 604     return Config.is_flag_already_in_argv(var)
 605
 606
 607 def print_usage() -> None:
 608     """Prints the normal help usage message out."""
 609     Config.print_usage()
 610
 611
 612 def usage() -> str:
 613     """
 614     Returns:
 615         program usage help text as a string.
 616     """
 617     return Config.usage()