3 # © Copyright 2021-2022, Scott Gasch
5 """Global configuration driven by commandline arguments, environment variables
6 and saved configuration files. This works across several modules.
14 parser = config.add_commandline_args(
16 "Args related to module doing the thing.",
19 "--module_do_the_thing",
22 help="Should the module do the thing?"
30 parser = config.add_commandline_args(
32 "A program that does the thing.",
38 help="Should we really do the thing?"
40 config.parse() # Very important, this must be invoked!
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`
54 if __name__ == '__main__':
57 Either way, you'll get this behavior from the commandline::
61 [--module_do_the_thing MODULE_DO_THE_THING]
65 Args related to module doing the thing.
67 --module_do_the_thing MODULE_DO_THE_THING
68 Should the module do the thing?
71 A program that does the thing
74 Should we really do the thing?
76 Arguments themselves should be accessed via
77 :code:`config.config['arg_name']`. e.g.::
79 if not config.config['dry_run']:
89 from typing import Any, Dict, List, Optional
91 from kazoo.client import KazooClient
95 # This module is commonly used by others in here and should avoid
96 # taking any unnecessary dependencies back on them.
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()
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
109 Use this, for example, when you need the helptext of an argument
110 to have its spacing preserved exactly, e.g.::
116 choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
118 help='''RAW|Our mode of operation. One of:
120 PLAY = play wordle with me! Pick a random solution or
121 specify a solution with --template.
123 CHEAT = given a --template and, optionally, --letters_in_word
124 and/or --letters_to_avoid, return the best guess word;
126 AUTOPLAY = given a complete word in --template, guess it step
127 by step showing work;
129 SELFTEST = autoplay every possible solution keeping track of
130 wins/losses and average number of guesses;
132 PRECOMPUTE = populate hash table with optimal guesses.
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)
143 # A global parser that we will collect arguments into.
144 ARGS = argparse.ArgumentParser(
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',
155 GROUP = ARGS.add_argument_group(
156 f'Global Config ({__file__})',
157 'Args that control the global config itself; how meta!',
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.',
169 help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
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.',
179 '--config_rejects_unrecognized_arguments',
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.',
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:
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.
200 # Has our parse() method been invoked yet?
201 self.config_parse_called = False
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] = {}
208 # Defer logging messages until later when logging has been
210 self.saved_messages: List[str] = []
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
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)
221 def __setitem__(self, key: str, value: Any) -> None:
222 self.config[key] = value
224 def __contains__(self, key: str) -> bool:
225 return key in self.config
228 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
229 """Create a new context for arguments and return a handle.
232 title: A title for your module's commandline arguments group.
233 description: A helpful description of your module.
236 An argparse._ArgumentGroup to be populated by the caller.
238 return ARGS.add_argument_group(title, description)
241 def overwrite_argparse_epilog(msg: str) -> None:
242 """Allows your code to override the default epilog created by
246 msg: The epilog message to substitute for the default.
251 def is_flag_already_in_argv(var: str) -> bool:
252 """Returns true if a particular flag is passed on the commandline
256 var: The flag to search for.
264 def print_usage() -> None:
265 """Prints the normal help usage message out."""
272 program usage help text as a string.
274 return ARGS.format_usage()
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
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)
290 reordered_action_groups.insert(0, grp)
291 return reordered_action_groups
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:
297 :code:`--argument_to_match`
301 :code:`ARGUMENT_TO_MATCH`
303 This allows programmers to set args via shell environment variables
304 in lieu of passing them on the cmdline.
307 usage_message = Config.usage()
311 # Foreach valid optional commandline option (line) generate
312 # its analogous environment variable.
313 for line in usage_message.split():
327 # Environment vars the same as flag names without
328 # the initial -'s and in UPPERCASE.
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}'
340 from string_utils import to_bool
342 if len(chunks) == 1 and to_bool(value):
344 elif len(chunks) > 1:
346 sys.argv.append(value)
351 def _augment_sys_argv_from_loadfile(self):
352 """Internal. Augment with arguments persisted in a saved file."""
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('=')
367 saw_other_args = True
369 if loadfile is not None:
371 if loadfile[:3] == 'zk:':
374 self.zk = KazooClient(
375 hosts=scott_secrets.ZOOKEEPER_NODES,
378 keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
379 keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
380 certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
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):
389 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
391 except Exception as e:
393 f'ERROR: Error talking with zookeeper while looking for {loadfile}'
395 elif not os.path.exists(loadfile):
397 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
401 msg = f'Augmenting commandline arguments with those from {loadfile}.'
403 msg = f'Reading commandline arguments from {loadfile}.'
404 print(msg, file=sys.stderr)
405 self.saved_messages.append(msg)
411 contents = self.zk.get(zkpath)[0]
412 contents = contents.decode()
415 for arg in contents.split('\n')
416 if 'config_savefile' not in arg
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}')
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]
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)
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.
442 if self.config_parse_called:
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.
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)
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()
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()
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))
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
475 if config['config_rejects_unrecognized_arguments']:
477 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
479 self.saved_messages.append(
480 f'Config encountered unrecognized commandline arguments: {unknown}'
482 sys.argv = sys.argv[:1] + unknown
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)
495 self.zk = KazooClient(
496 hosts=scott_secrets.ZOOKEEPER_NODES,
499 keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
500 keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
501 certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
504 if not self.zk.exists(zkpath):
505 self.zk.create(zkpath, data.encode())
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}')
512 with open(savefile, 'w') as wf:
515 # Also dump the config on stderr if requested.
516 if config['config_dump']:
519 self.config_parse_called = True
522 def has_been_parsed(self) -> bool:
523 """Returns True iff the global config has already been parsed"""
524 return self.config_parse_called
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:
534 # A global singleton instance of the Config class.
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.
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.
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.
553 title: A title for your module's commandline arguments group.
554 description: A helpful description of your module.
557 An argparse._ArgumentGroup to be populated by the caller.
559 return CONFIG.add_commandline_args(title, description)
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.
569 return CONFIG.parse(entry_module)
572 def has_been_parsed() -> bool:
573 """Returns True iff the global config has already been parsed"""
574 return CONFIG.has_been_parsed()
577 def late_logging() -> None:
578 """Log messages saved earlier now that logging has been initialized."""
579 CONFIG.late_logging()
582 def dump_config() -> None:
583 """Print the current config to stdout."""
587 def overwrite_argparse_epilog(msg: str) -> None:
588 """Allows your code to override the default epilog created by
592 msg: The epilog message to substitute for the default.
594 Config.overwrite_argparse_epilog(msg)
597 def is_flag_already_in_argv(var: str) -> bool:
598 """Returns true if a particular flag is passed on the commandline
602 var: The flag to search for.
604 return Config.is_flag_already_in_argv(var)
607 def print_usage() -> None:
608 """Prints the normal help usage message out."""
615 program usage help text as a string.
617 return Config.usage()