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