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']:
90 from typing import Any, Dict, List, Optional
92 from kazoo.client import KazooClient
96 # This module is commonly used by others in here and should avoid
97 # taking any unnecessary dependencies back on them.
99 # Defer logging messages until later when logging has been initialized.
100 SAVED_MESSAGES: List[str] = []
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()
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
113 Use this, for example, when you need the helptext of an argument
114 to have its spacing preserved exactly, e.g.::
120 choices=['CHEAT', 'AUTOPLAY', 'SELFTEST', 'PRECOMPUTE', 'PLAY'],
122 help='''RAW|Our mode of operation. One of:
124 PLAY = play wordle with me! Pick a random solution or
125 specify a solution with --template.
127 CHEAT = given a --template and, optionally, --letters_in_word
128 and/or --letters_to_avoid, return the best guess word;
130 AUTOPLAY = given a complete word in --template, guess it step
131 by step showing work;
133 SELFTEST = autoplay every possible solution keeping track of
134 wins/losses and average number of guesses;
136 PRECOMPUTE = populate hash table with optimal guesses.
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)
148 # A global parser that we will collect arguments into.
149 ARGS = argparse.ArgumentParser(
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',
161 # Keep track of if we've been called and prevent being called more
163 CONFIG_PARSE_CALLED = False
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] = {}
172 # It would be really nice if this shit worked from interactive python
175 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
176 """Create a new context for arguments and return a handle.
179 title: A title for your module's commandline arguments group.
180 description: A helpful description of your module.
183 An argparse._ArgumentGroup to be populated by the caller.
185 return ARGS.add_argument_group(title, description)
188 group = add_commandline_args(
189 f'Global Config ({__file__})',
190 'Args that control the global config itself; how meta!',
196 help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
202 help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
209 help='Populate config file compatible with --config_loadfile to save global config for later use.',
212 '--config_rejects_unrecognized_arguments',
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.'
223 def overwrite_argparse_epilog(msg: str) -> None:
224 """Allows your code to override the default epilog created by
228 msg: The epilog message to substitute for the default.
233 def is_flag_already_in_argv(var: str) -> bool:
234 """Returns true if a particular flag is passed on the commandline?
237 var: The flag to search for.
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
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)
258 reordered_action_groups.insert(0, grp)
259 return reordered_action_groups
262 def print_usage() -> None:
263 """Prints the normal help usage message out."""
270 program usage help text as a string.
272 return ARGS.format_usage()
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:
279 :code:`--argument_to_match`
283 :code:`ARGUMENT_TO_MATCH`
285 This allows programmers to set args via shell environment variables
286 in lieu of passing them on the cmdline.
290 usage_message = usage()
293 for x in usage_message.split():
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
316 if len(chunks) == 1 and to_bool(value):
318 elif len(chunks) > 1:
320 sys.argv.append(value)
325 def _augment_sys_argv_from_loadfile():
326 """Internal. Augment with arguments persisted in a saved file."""
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('=')
341 saw_other_args = True
343 if loadfile is not None:
345 if loadfile[:3] == 'zk:':
348 hosts=scott_secrets.ZOOKEEPER_NODES,
351 keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
352 keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
353 certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
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):
362 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
364 except Exception as e:
366 f'ERROR: Error talking with zookeeper while looking for {loadfile}'
368 elif not os.path.exists(loadfile):
370 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
374 msg = f'Augmenting commandline arguments with those from {loadfile}.'
376 msg = f'Reading commandline arguments from {loadfile}.'
377 print(msg, file=sys.stderr)
378 SAVED_MESSAGES.append(msg)
383 contents = zk.get(zkpath)[0]
384 contents = contents.decode()
386 arg.strip('\n') for arg in contents.split('\n') if 'config_savefile' not in arg
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}')
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]
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.
404 global CONFIG_PARSE_CALLED
405 if CONFIG_PARSE_CALLED:
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.
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)
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()
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()
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))
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
438 if config['config_rejects_unrecognized_arguments']:
440 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
442 SAVED_MESSAGES.append(f'Config encountered unrecognized commandline arguments: {unknown}')
443 sys.argv = sys.argv[:1] + unknown
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)
456 hosts=scott_secrets.ZOOKEEPER_NODES,
459 keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
460 keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
461 certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
464 if zk.exists(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}')
471 with open(savefile, 'w') as wf:
474 # Also dump the config on stderr if requested.
475 if config['config_dump']:
478 CONFIG_PARSE_CALLED = True
482 def has_been_parsed() -> bool:
483 """Returns True iff the global config has already been parsed"""
484 return CONFIG_PARSE_CALLED
488 """Print the current config to stdout."""
489 print("Global Configuration:", file=sys.stderr)
490 pprint.pprint(config, stream=sys.stderr)
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: