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 original_argv: List[str] = [arg for arg in sys.argv]
  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 # It would be really nice if this shit worked from interactive python
 120
 121
 122 def add_commandline_args(title: str, description: str = ""):
 123     """Create a new context for arguments and return a handle."""
 124     return args.add_argument_group(title, description)
 125
 126
 127 group = add_commandline_args(
 128     f'Global Config ({__file__})',
 129     'Args that control the global config itself; how meta!',
 130 )
 131 group.add_argument(
 132     '--config_loadfile',
 133     metavar='FILENAME',
 134     default=None,
 135     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 136 )
 137 group.add_argument(
 138     '--config_dump',
 139     default=False,
 140     action='store_true',
 141     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 142 )
 143 group.add_argument(
 144     '--config_savefile',
 145     type=str,
 146     metavar='FILENAME',
 147     default=None,
 148     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 149 )
 150 group.add_argument(
 151     '--config_rejects_unrecognized_arguments',
 152     default=False,
 153     action='store_true',
 154     help=(
 155         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 156         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 157         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 158     ),
 159 )
 160
 161
 162 def is_flag_already_in_argv(var: str):
 163     """Is a particular flag passed on the commandline?"""
 164     for _ in sys.argv:
 165         if var in _:
 166             return True
 167     return False
 168
 169
 170 def reorder_arg_action_groups(entry_module: Optional[str]):
 171     global program_name, args
 172     reordered_action_groups = []
 173     for group in args._action_groups:
 174         if entry_module is not None and entry_module in group.title:  # type: ignore
 175             reordered_action_groups.append(group)
 176         elif program_name in group.title:  # type: ignore
 177             reordered_action_groups.append(group)
 178         else:
 179             reordered_action_groups.insert(0, group)
 180     return reordered_action_groups
 181
 182
 183 def augment_sys_argv_from_environment_variables():
 184     global saved_messages
 185     usage_message = args.format_usage()
 186     optional = False
 187     var = ''
 188     for x in usage_message.split():
 189         if x[0] == '[':
 190             optional = True
 191         if optional:
 192             var += f'{x} '
 193             if x[-1] == ']':
 194                 optional = False
 195                 var = var.strip()
 196                 var = var.strip('[')
 197                 var = var.strip(']')
 198                 chunks = var.split()
 199                 if len(chunks) > 1:
 200                     var = var.split()[0]
 201
 202                 # Environment vars the same as flag names without
 203                 # the initial -'s and in UPPERCASE.
 204                 env = var.strip('-').upper()
 205                 if env in os.environ:
 206                     if not is_flag_already_in_argv(var):
 207                         value = os.environ[env]
 208                         saved_messages.append(f'Initialized from environment: {var} = {value}')
 209                         from string_utils import to_bool
 210
 211                         if len(chunks) == 1 and to_bool(value):
 212                             sys.argv.append(var)
 213                         elif len(chunks) > 1:
 214                             sys.argv.append(var)
 215                             sys.argv.append(value)
 216                 var = ''
 217                 env = ''
 218
 219
 220 def augment_sys_argv_from_loadfile():
 221     global saved_messages
 222     loadfile = None
 223     saw_other_args = False
 224     grab_next_arg = False
 225     for arg in sys.argv[1:]:
 226         if 'config_loadfile' in arg:
 227             pieces = arg.split('=')
 228             if len(pieces) > 1:
 229                 loadfile = pieces[1]
 230             else:
 231                 grab_next_arg = True
 232         elif grab_next_arg:
 233             loadfile = arg
 234         else:
 235             saw_other_args = True
 236
 237     if loadfile is not None:
 238         if not os.path.exists(loadfile):
 239             raise Exception(
 240                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 241             )
 242         if saw_other_args:
 243             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 244         else:
 245             msg = f'Reading commandline arguments from {loadfile}.'
 246         print(msg, file=sys.stderr)
 247         saved_messages.append(msg)
 248
 249         with open(loadfile, 'r') as rf:
 250             newargs = rf.readlines()
 251         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 252         sys.argv += newargs
 253
 254
 255 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 256     """Main program should call this early in main().  Note that the
 257     bootstrap.initialize wrapper takes care of this automatically.
 258
 259     """
 260     global config_parse_called
 261     if config_parse_called:
 262         return config
 263     global saved_messages
 264
 265     # If we're about to do the usage message dump, put the main
 266     # module's argument group last in the list (if possible) so that
 267     # when the user passes -h or --help, it will be visible on the
 268     # screen w/o scrolling.
 269     for arg in sys.argv:
 270         if arg == '--help' or arg == '-h':
 271             args._action_groups = reorder_arg_action_groups(entry_module)
 272
 273     # Examine the environment for variables that match known flags.
 274     # For a flag called --example_flag the corresponding environment
 275     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 276     # these into sys.argv to be parsed.
 277     augment_sys_argv_from_environment_variables()
 278
 279     # Look for loadfile and read/parse it if present.  This also
 280     # works by jamming these values onto sys.argv.
 281     augment_sys_argv_from_loadfile()
 282
 283     # Parse (possibly augmented, possibly completely overwritten)
 284     # commandline args with argparse normally and populate config.
 285     known, unknown = args.parse_known_args()
 286     config.update(vars(known))
 287
 288     # Reconstruct the argv with unrecognized flags for the benefit of
 289     # future argument parsers.  For example, unittest_main in python
 290     # has some of its own flags.  If we didn't recognize it, maybe
 291     # someone else will.
 292     if len(unknown) > 0:
 293         if config['config_rejects_unrecognized_arguments']:
 294             raise Exception(
 295                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 296             )
 297         saved_messages.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 298     sys.argv = sys.argv[:1] + unknown
 299
 300     # Check for savefile and populate it if requested.
 301     savefile = config['config_savefile']
 302     if savefile and len(savefile) > 0:
 303         with open(savefile, 'w') as wf:
 304             wf.write("\n".join(original_argv[1:]))
 305
 306     # Also dump the config on stderr if requested.
 307     if config['config_dump']:
 308         dump_config()
 309
 310     config_parse_called = True
 311     return config
 312
 313
 314 def has_been_parsed() -> bool:
 315     """Has the global config been parsed yet?"""
 316     global config_parse_called
 317     return config_parse_called
 318
 319
 320 def dump_config():
 321     """Print the current config to stdout."""
 322     print("Global Configuration:", file=sys.stderr)
 323     pprint.pprint(config, stream=sys.stderr)
 324     print()
 325
 326
 327 def late_logging():
 328     """Log messages saved earlier now that logging has been initialized."""
 329     logger = logging.getLogger(__name__)
 330     global saved_messages
 331     for _ in saved_messages:
 332         logger.debug(_)