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