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 and saved configuration files.  This works across several modules.
   7
   8 Usage:
   9
  10     module.py:
  11     ----------
  12     import config
  13
  14     parser = config.add_commandline_args(
  15         "Module",
  16         "Args related to module doing the thing.",
  17     )
  18     parser.add_argument(
  19         "--module_do_the_thing",
  20         type=bool,
  21         default=True,
  22         help="Should the module do the thing?"
  23     )
  24
  25     main.py:
  26     --------
  27     import config
  28
  29     def main() -> None:
  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         config.parse()   # Very important, this must be invoked!
  41
  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 bootstrap module's initialize wrapper.
  45
  46     % main.py -h
  47     usage: main.py [-h]
  48                    [--module_do_the_thing MODULE_DO_THE_THING]
  49                    [--dry_run DRY_RUN]
  50
  51     Module:
  52       Args related to module doing the thing.
  53
  54       --module_do_the_thing MODULE_DO_THE_THING
  55                    Should the module do the thing?
  56
  57     Main:
  58       A program that does the thing
  59
  60       --dry_run
  61                    Should we really do the thing?
  62
  63     Arguments themselves should be accessed via
  64     config.config['arg_name'].  e.g.
  65
  66     if not config.config['dry_run']:
  67         module.do_the_thing()
  68
  69 """
  70
  71 import argparse
  72 import logging
  73 import os
  74 import pprint
  75 import sys
  76 from typing import Any, Dict, List, Optional
  77
  78 # This module is commonly used by others in here and should avoid
  79 # taking any unnecessary dependencies back on them.
  80
  81 # Defer logging messages until later when logging has been initialized.
  82 SAVED_MESSAGES: List[str] = []
  83
  84 # Make a copy of the original program arguments.
  85 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
  86 ORIG_ARGV: List[str] = sys.argv.copy()
  87
  88
  89 class OptionalRawFormatter(argparse.HelpFormatter):
  90     """This formatter has the same bahavior as the normal argparse text
  91     formatter except when the help text of an argument begins with
  92     "RAW|".  In that case, the line breaks are preserved and the text
  93     is not wrapped.
  94
  95     """
  96
  97     def _split_lines(self, text, width):
  98         if text.startswith('RAW|'):
  99             return text[4:].splitlines()
 100         return argparse.HelpFormatter._split_lines(self, text, width)
 101
 102
 103 # A global parser that we will collect arguments into.
 104 ARGS = argparse.ArgumentParser(
 105     description=None,
 106     formatter_class=OptionalRawFormatter,
 107     fromfile_prefix_chars="@",
 108     epilog=f'{PROGRAM_NAME} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
 109 )
 110
 111 # Keep track of if we've been called and prevent being called more
 112 # than once.
 113 CONFIG_PARSE_CALLED = False
 114
 115
 116 # A global configuration dictionary that will contain parsed arguments.
 117 # It is also this variable that modules use to access parsed arguments.
 118 # This is the data that is most interesting to our callers; it will hold
 119 # the configuration result.
 120 config: Dict[str, Any] = {}
 121
 122 # It would be really nice if this shit worked from interactive python
 123
 124
 125 def add_commandline_args(title: str, description: str = ""):
 126     """Create a new context for arguments and return a handle."""
 127     return ARGS.add_argument_group(title, description)
 128
 129
 130 group = add_commandline_args(
 131     f'Global Config ({__file__})',
 132     'Args that control the global config itself; how meta!',
 133 )
 134 group.add_argument(
 135     '--config_loadfile',
 136     metavar='FILENAME',
 137     default=None,
 138     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 139 )
 140 group.add_argument(
 141     '--config_dump',
 142     default=False,
 143     action='store_true',
 144     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 145 )
 146 group.add_argument(
 147     '--config_savefile',
 148     type=str,
 149     metavar='FILENAME',
 150     default=None,
 151     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 152 )
 153 group.add_argument(
 154     '--config_rejects_unrecognized_arguments',
 155     default=False,
 156     action='store_true',
 157     help=(
 158         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 159         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 160         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 161     ),
 162 )
 163
 164
 165 def overwrite_argparse_epilog(msg: str) -> None:
 166     ARGS.epilog = msg
 167
 168
 169 def is_flag_already_in_argv(var: str):
 170     """Is a particular flag passed on the commandline?"""
 171     for _ in sys.argv:
 172         if var in _:
 173             return True
 174     return False
 175
 176
 177 def reorder_arg_action_groups_before_help(entry_module: Optional[str]):
 178     reordered_action_groups = []
 179     for grp in ARGS._action_groups:
 180         if entry_module is not None and entry_module in grp.title:  # type: ignore
 181             reordered_action_groups.append(grp)
 182         elif PROGRAM_NAME in group.title:  # type: ignore
 183             reordered_action_groups.append(grp)
 184         else:
 185             reordered_action_groups.insert(0, grp)
 186     return reordered_action_groups
 187
 188
 189 def augment_sys_argv_from_environment_variables():
 190     usage_message = ARGS.format_usage()
 191     optional = False
 192     var = ''
 193     for x in usage_message.split():
 194         if x[0] == '[':
 195             optional = True
 196         if optional:
 197             var += f'{x} '
 198             if x[-1] == ']':
 199                 optional = False
 200                 var = var.strip()
 201                 var = var.strip('[')
 202                 var = var.strip(']')
 203                 chunks = var.split()
 204                 if len(chunks) > 1:
 205                     var = var.split()[0]
 206
 207                 # Environment vars the same as flag names without
 208                 # the initial -'s and in UPPERCASE.
 209                 env = var.strip('-').upper()
 210                 if env in os.environ:
 211                     if not is_flag_already_in_argv(var):
 212                         value = os.environ[env]
 213                         SAVED_MESSAGES.append(f'Initialized from environment: {var} = {value}')
 214                         from string_utils import to_bool
 215
 216                         if len(chunks) == 1 and to_bool(value):
 217                             sys.argv.append(var)
 218                         elif len(chunks) > 1:
 219                             sys.argv.append(var)
 220                             sys.argv.append(value)
 221                 var = ''
 222                 env = ''
 223
 224
 225 def augment_sys_argv_from_loadfile():
 226     loadfile = None
 227     saw_other_args = False
 228     grab_next_arg = False
 229     for arg in sys.argv[1:]:
 230         if 'config_loadfile' in arg:
 231             pieces = arg.split('=')
 232             if len(pieces) > 1:
 233                 loadfile = pieces[1]
 234             else:
 235                 grab_next_arg = True
 236         elif grab_next_arg:
 237             loadfile = arg
 238         else:
 239             saw_other_args = True
 240
 241     if loadfile is not None:
 242         if not os.path.exists(loadfile):
 243             raise Exception(
 244                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 245             )
 246         if saw_other_args:
 247             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 248         else:
 249             msg = f'Reading commandline arguments from {loadfile}.'
 250         print(msg, file=sys.stderr)
 251         SAVED_MESSAGES.append(msg)
 252
 253         with open(loadfile, 'r') as rf:
 254             newargs = rf.readlines()
 255         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 256         sys.argv += newargs
 257
 258
 259 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 260     """Main program should call this early in main().  Note that the
 261     bootstrap.initialize wrapper takes care of this automatically.
 262
 263     """
 264     global CONFIG_PARSE_CALLED
 265     if CONFIG_PARSE_CALLED:
 266         return config
 267
 268     # If we're about to do the usage message dump, put the main
 269     # module's argument group last in the list (if possible) so that
 270     # when the user passes -h or --help, it will be visible on the
 271     # screen w/o scrolling.
 272     for arg in sys.argv:
 273         if arg in ('--help', '-h'):
 274             if entry_module is not None:
 275                 entry_module = os.path.basename(entry_module)
 276             ARGS._action_groups = reorder_arg_action_groups_before_help(entry_module)
 277
 278     # Examine the environment for variables that match known flags.
 279     # For a flag called --example_flag the corresponding environment
 280     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 281     # these into sys.argv to be parsed.
 282     augment_sys_argv_from_environment_variables()
 283
 284     # Look for loadfile and read/parse it if present.  This also
 285     # works by jamming these values onto sys.argv.
 286     augment_sys_argv_from_loadfile()
 287
 288     # Parse (possibly augmented, possibly completely overwritten)
 289     # commandline args with argparse normally and populate config.
 290     known, unknown = ARGS.parse_known_args()
 291     config.update(vars(known))
 292
 293     # Reconstruct the argv with unrecognized flags for the benefit of
 294     # future argument parsers.  For example, unittest_main in python
 295     # has some of its own flags.  If we didn't recognize it, maybe
 296     # someone else will.
 297     if len(unknown) > 0:
 298         if config['config_rejects_unrecognized_arguments']:
 299             raise Exception(
 300                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 301             )
 302         SAVED_MESSAGES.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 303     sys.argv = sys.argv[:1] + unknown
 304
 305     # Check for savefile and populate it if requested.
 306     savefile = config['config_savefile']
 307     if savefile and len(savefile) > 0:
 308         with open(savefile, 'w') as wf:
 309             wf.write("\n".join(ORIG_ARGV[1:]))
 310
 311     # Also dump the config on stderr if requested.
 312     if config['config_dump']:
 313         dump_config()
 314
 315     CONFIG_PARSE_CALLED = True
 316     return config
 317
 318
 319 def has_been_parsed() -> bool:
 320     """Has the global config been parsed yet?"""
 321     return CONFIG_PARSE_CALLED
 322
 323
 324 def dump_config():
 325     """Print the current config to stdout."""
 326     print("Global Configuration:", file=sys.stderr)
 327     pprint.pprint(config, stream=sys.stderr)
 328     print()
 329
 330
 331 def late_logging():
 332     """Log messages saved earlier now that logging has been initialized."""
 333     logger = logging.getLogger(__name__)
 334     logger.debug('Original commandline was: %s', ORIG_ARGV)
 335     for _ in SAVED_MESSAGES:
 336         logger.debug(_)