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 = os.path.basename(sys.argv[0])
  84 original_argv = [arg for arg in sys.argv]
  85
  86
  87 # A global parser that we will collect arguments into.
  88 args = argparse.ArgumentParser(
  89     description=None,
  90     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  91     fromfile_prefix_chars="@",
  92     epilog=f'{program_name} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.'
  93 )
  94
  95 # Keep track of if we've been called and prevent being called more
  96 # than once.
  97 config_parse_called = False
  98
  99
 100 # A global configuration dictionary that will contain parsed arguments.
 101 # It is also this variable that modules use to access parsed arguments.
 102 # This is the data that is most interesting to our callers; it will hold
 103 # the configuration result.
 104 config = {}
 105 # It would be really nice if this shit worked from interactive python
 106
 107
 108 def add_commandline_args(title: str, description: str = ""):
 109     """Create a new context for arguments and return a handle."""
 110     return args.add_argument_group(title, description)
 111
 112
 113 group = add_commandline_args(
 114     f'Global Config ({__file__})',
 115     'Args that control the global config itself; how meta!',
 116 )
 117 group.add_argument(
 118     '--config_loadfile',
 119     metavar='FILENAME',
 120     default=None,
 121     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 122 )
 123 group.add_argument(
 124     '--config_dump',
 125     default=False,
 126     action='store_true',
 127     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 128 )
 129 group.add_argument(
 130     '--config_savefile',
 131     type=str,
 132     metavar='FILENAME',
 133     default=None,
 134     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 135 )
 136
 137
 138 def is_flag_already_in_argv(var: str):
 139     """Is a particular flag passed on the commandline?"""
 140     for _ in sys.argv:
 141         if var in _:
 142             return True
 143     return False
 144
 145
 146 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 147     """Main program should call this early in main().  Note that the
 148     bootstrap.initialize wrapper takes care of this automatically.
 149
 150     """
 151     global config_parse_called
 152     if config_parse_called:
 153         return config
 154
 155     global saved_messages
 156
 157     # If we're about to do the usage message dump, put the main module's
 158     # argument group last in the list (if possible) so that when the user
 159     # passes -h or --help, it will be visible on the screen w/o scrolling.
 160     reordered_action_groups = []
 161     global prog
 162     for arg in sys.argv:
 163         if arg == '--help' or arg == '-h':
 164             for group in args._action_groups:
 165                 if entry_module is not None and entry_module in group.title:
 166                     reordered_action_groups.append(group)
 167                 elif program_name in group.title:
 168                     reordered_action_groups.append(group)
 169                 else:
 170                     reordered_action_groups.insert(0, group)
 171             args._action_groups = reordered_action_groups
 172
 173     # Examine the environment for variables that match known flags.
 174     # For a flag called --example_flag the corresponding environment
 175     # variable would be called EXAMPLE_FLAG.
 176     usage_message = args.format_usage()
 177     optional = False
 178     var = ''
 179     for x in usage_message.split():
 180         if x[0] == '[':
 181             optional = True
 182         if optional:
 183             var += f'{x} '
 184             if x[-1] == ']':
 185                 optional = False
 186                 var = var.strip()
 187                 var = var.strip('[')
 188                 var = var.strip(']')
 189                 chunks = var.split()
 190                 if len(chunks) > 1:
 191                     var = var.split()[0]
 192
 193                 # Environment vars the same as flag names without
 194                 # the initial -'s and in UPPERCASE.
 195                 env = var.strip('-').upper()
 196                 if env in os.environ:
 197                     if not is_flag_already_in_argv(var):
 198                         value = os.environ[env]
 199                         saved_messages.append(
 200                             f'Initialized from environment: {var} = {value}'
 201                         )
 202                         from string_utils import to_bool
 203                         if len(chunks) == 1 and to_bool(value):
 204                             sys.argv.append(var)
 205                         elif len(chunks) > 1:
 206                             sys.argv.append(var)
 207                             sys.argv.append(value)
 208                 var = ''
 209                 env = ''
 210         else:
 211             next
 212
 213     # Look for loadfile and read/parse it if present.
 214     loadfile = None
 215     saw_other_args = False
 216     grab_next_arg = False
 217     for arg in sys.argv[1:]:
 218         if 'config_loadfile' in arg:
 219             pieces = arg.split('=')
 220             if len(pieces) > 1:
 221                 loadfile = pieces[1]
 222             else:
 223                 grab_next_arg = True
 224         elif grab_next_arg:
 225             loadfile = arg
 226         else:
 227             saw_other_args = True
 228
 229     if loadfile is not None:
 230         if saw_other_args:
 231             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 232             print(msg, file=sys.stderr)
 233             saved_messages.append(msg)
 234         if not os.path.exists(loadfile):
 235             print(f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.',
 236                   file=sys.stderr)
 237             sys.exit(-1)
 238         with open(loadfile, 'r') as rf:
 239             newargs = rf.readlines()
 240         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 241         sys.argv += newargs
 242
 243     # Parse (possibly augmented, possibly completely overwritten)
 244     # commandline args with argparse normally and populate config.
 245     known, unknown = args.parse_known_args()
 246     config.update(vars(known))
 247
 248     # Reconstruct the argv with unrecognized flags for the benefit of
 249     # future argument parsers.  For example, unittest_main in python
 250     # has some of its own flags.  If we didn't recognize it, maybe
 251     # someone else will.
 252     sys.argv = sys.argv[:1] + unknown
 253
 254     # Check for savefile and populate it if requested.
 255     savefile = config['config_savefile']
 256     if savefile and len(savefile) > 0:
 257         with open(savefile, 'w') as wf:
 258             wf.write(
 259                 "\n".join(original_argv[1:])
 260             )
 261
 262     # Also dump the config on stderr if requested.
 263     if config['config_dump']:
 264         dump_config()
 265
 266     config_parse_called = True
 267     return config
 268
 269
 270 def has_been_parsed() -> bool:
 271     """Has the global config been parsed yet?"""
 272     global config_parse_called
 273     return config_parse_called
 274
 275
 276 def dump_config():
 277     """Print the current config to stdout."""
 278     print("Global Configuration:", file=sys.stderr)
 279     pprint.pprint(config, stream=sys.stderr)
 280     print()
 281
 282
 283 def late_logging():
 284     """Log messages saved earlier now that logging has been initialized."""
 285     logger = logging.getLogger(__name__)
 286     global saved_messages
 287     for _ in saved_messages:
 288         logger.debug(_)