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