config.py

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