More new config bugs: needs to support a .get() method... doh.
[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 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     def get(self, key: str, default: Any = None) -> Optional[Any]:
238         return self.config.get(key, default)
239
240     @staticmethod
241     def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
242         """Create a new context for arguments and return a handle.
243
244         Args:
245             title: A title for your module's commandline arguments group.
246             description: A helpful description of your module.
247
248         Returns:
249             An argparse._ArgumentGroup to be populated by the caller.
250         """
251         return ARGS.add_argument_group(title, description)
252
253     @staticmethod
254     def overwrite_argparse_epilog(msg: str) -> None:
255         """Allows your code to override the default epilog created by
256         argparse.
257
258         Args:
259             msg: The epilog message to substitute for the default.
260         """
261         ARGS.epilog = msg
262
263     @staticmethod
264     def is_flag_already_in_argv(var: str) -> bool:
265         """Returns true if a particular flag is passed on the commandline
266         and false otherwise.
267
268         Args:
269             var: The flag to search for.
270         """
271         for _ in sys.argv:
272             if var in _:
273                 return True
274         return False
275
276     @staticmethod
277     def print_usage() -> None:
278         """Prints the normal help usage message out."""
279         ARGS.print_help()
280
281     @staticmethod
282     def usage() -> str:
283         """
284         Returns:
285             program usage help text as a string.
286         """
287         return ARGS.format_usage()
288
289     @staticmethod
290     def _reorder_arg_action_groups_before_help(entry_module: Optional[str]):
291         """Internal.  Used to reorder the arguments before dumping out a
292         generated help string such that the main program's arguments come
293         last.
294
295         """
296         reordered_action_groups = []
297         for grp in ARGS._action_groups:
298             if entry_module is not None and entry_module in grp.title:  # type: ignore
299                 reordered_action_groups.append(grp)
300             elif PROGRAM_NAME in GROUP.title:  # type: ignore
301                 reordered_action_groups.append(grp)
302             else:
303                 reordered_action_groups.insert(0, grp)
304         return reordered_action_groups
305
306     @staticmethod
307     def _parse_arg_into_env(arg: str) -> Optional[Tuple[str, str, List[str]]]:
308         """Internal helper to parse commandling args into environment vars."""
309         arg = arg.strip()
310         if not arg.startswith('['):
311             return None
312         arg = arg.strip('[')
313         if not arg.endswith(']'):
314             return None
315         arg = arg.strip(']')
316
317         chunks = arg.split()
318         if len(chunks) > 1:
319             var = chunks[0]
320         else:
321             var = arg
322
323         # Environment vars the same as flag names without
324         # the initial -'s and in UPPERCASE.
325         env = var.upper()
326         while env[0] == '-':
327             env = env[1:]
328         return var, env, chunks
329
330     def _augment_sys_argv_from_environment_variables(self):
331         """Internal.  Look at the system environment for variables that match
332         commandline arg names.  This is done via some munging such that:
333
334         :code:`--argument_to_match`
335
336         ...is matched by:
337
338         :code:`ARGUMENT_TO_MATCH`
339
340         This allows users to set args via shell environment variables
341         in lieu of passing them on the cmdline.
342
343         """
344         usage_message = Config.usage()
345         optional = False
346         arg = ''
347
348         # Foreach valid optional commandline option (chunk) generate
349         # its analogous environment variable.
350         for chunk in usage_message.split():
351             if chunk[0] == '[':
352                 optional = True
353             if optional:
354                 arg += f'{chunk} '
355                 if chunk[-1] == ']':
356                     optional = False
357                     _ = Config._parse_arg_into_env(arg)
358                     if _:
359                         var, env, chunks = _
360                         if env in os.environ:
361                             if not Config.is_flag_already_in_argv(var):
362                                 value = os.environ[env]
363                                 self.saved_messages.append(
364                                     f'Initialized from environment: {var} = {value}'
365                                 )
366                                 from string_utils import to_bool
367
368                                 if len(chunks) == 1 and to_bool(value):
369                                     sys.argv.append(var)
370                                 elif len(chunks) > 1:
371                                     sys.argv.append(var)
372                                     sys.argv.append(value)
373                     arg = ''
374
375     def _process_dynamic_args(self, event: WatchedEvent):
376         assert self.zk
377         logger = logging.getLogger(__name__)
378         contents, meta = self.zk.get(event.path, watch=self._process_dynamic_args)
379         logger.debug('Update for %s at version=%d.', event.path, meta.version)
380         logger.debug(
381             'Max known version for %s is %d.', event.path, self.max_version.get(event.path, 0)
382         )
383         if meta.version > self.max_version.get(event.path, 0):
384             self.max_version[event.path] = meta.version
385             contents = contents.decode()
386             temp_argv = []
387             for arg in contents.split():
388                 if 'dynamic' in arg:
389                     temp_argv.append(arg)
390                     logger.info("Updating %s from zookeeper async config change.", arg)
391             if len(temp_argv) > 0:
392                 old_argv = sys.argv
393                 sys.argv = temp_argv
394                 known, _ = ARGS.parse_known_args()
395                 sys.argv = old_argv
396                 self.config.update(vars(known))
397
398     def _augment_sys_argv_from_loadfile(self):
399         """Internal.  Augment with arguments persisted in a saved file."""
400
401         loadfile = None
402         saw_other_args = False
403         grab_next_arg = False
404         for arg in sys.argv[1:]:
405             if 'config_loadfile' in arg:
406                 pieces = arg.split('=')
407                 if len(pieces) > 1:
408                     loadfile = pieces[1]
409                 else:
410                     grab_next_arg = True
411             elif grab_next_arg:
412                 loadfile = arg
413             else:
414                 saw_other_args = True
415
416         if loadfile is not None:
417             zkpath = None
418             if loadfile[:3] == 'zk:':
419                 try:
420                     if self.zk is None:
421                         self.zk = KazooClient(
422                             hosts=scott_secrets.ZOOKEEPER_NODES,
423                             use_ssl=True,
424                             verify_certs=False,
425                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
426                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
427                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
428                         )
429                         self.zk.start()
430                     zkpath = loadfile[3:]
431                     if not zkpath.startswith('/config/'):
432                         zkpath = '/config/' + zkpath
433                         zkpath = re.sub(r'//+', '/', zkpath)
434                     if not self.zk.exists(zkpath):
435                         raise Exception(
436                             f'ERROR: --config_loadfile argument must be a file, {loadfile} not found (in zookeeper)'
437                         )
438                 except Exception as e:
439                     raise Exception(
440                         f'ERROR: Error talking with zookeeper while looking for {loadfile}'
441                     ) from e
442             elif not os.path.exists(loadfile):
443                 raise Exception(
444                     f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
445                 )
446
447             if saw_other_args:
448                 msg = f'Augmenting commandline arguments with those from {loadfile}.'
449             else:
450                 msg = f'Reading commandline arguments from {loadfile}.'
451             print(msg, file=sys.stderr)
452             self.saved_messages.append(msg)
453
454             newargs = []
455             if zkpath:
456                 try:
457                     assert self.zk
458                     contents, meta = self.zk.get(zkpath, watch=self._process_dynamic_args)
459                     contents = contents.decode()
460                     newargs = [
461                         arg.strip('\n')
462                         for arg in contents.split('\n')
463                         if 'config_savefile' not in arg
464                     ]
465                     self.saved_messages.append(f'Setting {zkpath}\'s max_version to {meta.version}')
466                     self.max_version[zkpath] = meta.version
467                 except Exception as e:
468                     raise Exception(f'Error reading {zkpath} from zookeeper.') from e
469                 self.saved_messages.append(f'Loaded config from zookeeper from {zkpath}')
470             else:
471                 with open(loadfile, 'r') as rf:
472                     newargs = rf.readlines()
473                 newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
474             sys.argv += newargs
475
476     def dump_config(self):
477         """Print the current config to stdout."""
478         print("Global Configuration:", file=sys.stderr)
479         pprint.pprint(self.config, stream=sys.stderr)
480         print()
481
482     def parse(self, entry_module: Optional[str]) -> Dict[str, Any]:
483         """Main program should call this early in main().  Note that the
484         :code:`bootstrap.initialize` wrapper takes care of this automatically.
485         This should only be called once per program invocation.
486
487         """
488         if self.config_parse_called:
489             return self.config
490
491         # If we're about to do the usage message dump, put the main
492         # module's argument group last in the list (if possible) so that
493         # when the user passes -h or --help, it will be visible on the
494         # screen w/o scrolling.
495         for arg in sys.argv:
496             if arg in ('--help', '-h'):
497                 if entry_module is not None:
498                     entry_module = os.path.basename(entry_module)
499                 ARGS._action_groups = Config._reorder_arg_action_groups_before_help(entry_module)
500
501         # Examine the environment for variables that match known flags.
502         # For a flag called --example_flag the corresponding environment
503         # variable would be called EXAMPLE_FLAG.  If found, hackily add
504         # these into sys.argv to be parsed.
505         self._augment_sys_argv_from_environment_variables()
506
507         # Look for loadfile and read/parse it if present.  This also
508         # works by jamming these values onto sys.argv.
509         self._augment_sys_argv_from_loadfile()
510
511         # Parse (possibly augmented, possibly completely overwritten)
512         # commandline args with argparse normally and populate config.
513         known, unknown = ARGS.parse_known_args()
514         self.config.update(vars(known))
515
516         # Reconstruct the argv with unrecognized flags for the benefit of
517         # future argument parsers.  For example, unittest_main in python
518         # has some of its own flags.  If we didn't recognize it, maybe
519         # someone else will.
520         if len(unknown) > 0:
521             if config['config_rejects_unrecognized_arguments']:
522                 raise Exception(
523                     f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
524                 )
525             self.saved_messages.append(
526                 f'Config encountered unrecognized commandline arguments: {unknown}'
527             )
528         sys.argv = sys.argv[:1] + unknown
529
530         # Check for savefile and populate it if requested.
531         savefile = config['config_savefile']
532         if savefile and len(savefile) > 0:
533             data = '\n'.join(ORIG_ARGV[1:])
534             if savefile[:3] == 'zk:':
535                 zkpath = savefile[3:]
536                 if not zkpath.startswith('/config/'):
537                     zkpath = '/config/' + zkpath
538                     zkpath = re.sub(r'//+', '/', zkpath)
539                 try:
540                     if not self.zk:
541                         self.zk = KazooClient(
542                             hosts=scott_secrets.ZOOKEEPER_NODES,
543                             use_ssl=True,
544                             verify_certs=False,
545                             keyfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
546                             keyfile_password=scott_secrets.ZOOKEEPER_CLIENT_PASS,
547                             certfile=scott_secrets.ZOOKEEPER_CLIENT_CERT,
548                         )
549                         self.zk.start()
550                     data = data.encode()
551                     if len(data) > 1024 * 1024:
552                         raise Exception(f'Saved args are too large!  ({len(data)} bytes)')
553                     if not self.zk.exists(zkpath):
554                         self.zk.create(zkpath, data)
555                         self.saved_messages.append(
556                             f'Just created {zkpath}; setting its max_version to 0'
557                         )
558                         self.max_version[zkpath] = 0
559                     else:
560                         meta = self.zk.set(zkpath, data)
561                         self.saved_messages.append(
562                             f'Setting {zkpath}\'s max_version to {meta.version}'
563                         )
564                         self.max_version[zkpath] = meta.version
565                 except Exception as e:
566                     raise Exception(f'Failed to create zookeeper path {zkpath}') from e
567                 self.saved_messages.append(f'Saved config to zookeeper in {zkpath}')
568             else:
569                 with open(savefile, 'w') as wf:
570                     wf.write(data)
571
572         # Also dump the config on stderr if requested.
573         if config['config_dump']:
574             self.dump_config()
575
576         self.config_parse_called = True
577         return self.config
578
579     def has_been_parsed(self) -> bool:
580         """Returns True iff the global config has already been parsed"""
581         return self.config_parse_called
582
583     def late_logging(self):
584         """Log messages saved earlier now that logging has been initialized."""
585         logger = logging.getLogger(__name__)
586         logger.debug('Original commandline was: %s', ORIG_ARGV)
587         for _ in self.saved_messages:
588             logger.debug(_)
589
590
591 # A global singleton instance of the Config class.
592 CONFIG = Config()
593
594 # A lot of client code uses config.config['whatever'] to lookup
595 # configuration so to preserve this we make this, config.config, with
596 # a __getitem__ method on it.
597 config = CONFIG
598
599 # Config didn't use to be a class; it was a mess of module-level
600 # functions and data.  The functions below preserve the old interface
601 # so that existing clients do not need to be changed.  As you can see,
602 # they mostly just thunk into the config class.
603
604
605 def add_commandline_args(title: str, description: str = "") -> argparse._ArgumentGroup:
606     """Create a new context for arguments and return a handle.  An alias
607     for config.config.add_commandline_args.
608
609     Args:
610         title: A title for your module's commandline arguments group.
611         description: A helpful description of your module.
612
613     Returns:
614         An argparse._ArgumentGroup to be populated by the caller.
615     """
616     return CONFIG.add_commandline_args(title, description)
617
618
619 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
620     """Main program should call this early in main().  Note that the
621     :code:`bootstrap.initialize` wrapper takes care of this automatically.
622     This should only be called once per program invocation.  Subsequent
623     calls do not reparse the configuration settings but rather just
624     return the current state.
625     """
626     return CONFIG.parse(entry_module)
627
628
629 def has_been_parsed() -> bool:
630     """Returns True iff the global config has already been parsed"""
631     return CONFIG.has_been_parsed()
632
633
634 def late_logging() -> None:
635     """Log messages saved earlier now that logging has been initialized."""
636     CONFIG.late_logging()
637
638
639 def dump_config() -> None:
640     """Print the current config to stdout."""
641     CONFIG.dump_config()
642
643
644 def overwrite_argparse_epilog(msg: str) -> None:
645     """Allows your code to override the default epilog created by
646     argparse.
647
648     Args:
649         msg: The epilog message to substitute for the default.
650     """
651     Config.overwrite_argparse_epilog(msg)
652
653
654 def is_flag_already_in_argv(var: str) -> bool:
655     """Returns true if a particular flag is passed on the commandline
656     and false otherwise.
657
658     Args:
659         var: The flag to search for.
660     """
661     return Config.is_flag_already_in_argv(var)
662
663
664 def print_usage() -> None:
665     """Prints the normal help usage message out."""
666     Config.print_usage()
667
668
669 def usage() -> str:
670     """
671     Returns:
672         program usage help text as a string.
673     """
674     return Config.usage()