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'------------------------------------------------------------------------------\n{program_name} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.\n------------------------------------------------------------------------------'
  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 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 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()"""
 146     global config_parse_called
 147     if config_parse_called:
 148         return config
 149
 150     global saved_messages
 151
 152     # If we're about to do the usage message dump, put the main module's
 153     # argument group last in the list (if possible) so that when the user
 154     # passes -h or --help, it will be visible on the screen w/o scrolling.
 155     reordered_action_groups = []
 156     global prog
 157     for arg in sys.argv:
 158         if arg == '--help' or arg == '-h':
 159             for group in args._action_groups:
 160                 if entry_module is not None and entry_module in group.title:
 161                     reordered_action_groups.append(group)
 162                 elif program_name in group.title:
 163                     reordered_action_groups.append(group)
 164                 else:
 165                     reordered_action_groups.insert(0, group)
 166             args._action_groups = reordered_action_groups
 167
 168     # Examine the environment variables that match known flags.  For a
 169     # flag called --example_flag the corresponding environment
 170     # variable would be called EXAMPLE_FLAG.
 171     usage_message = args.format_usage()
 172     optional = False
 173     var = ''
 174     for x in usage_message.split():
 175         if x[0] == '[':
 176             optional = True
 177         if optional:
 178             var += f'{x} '
 179             if x[-1] == ']':
 180                 optional = False
 181                 var = var.strip()
 182                 var = var.strip('[')
 183                 var = var.strip(']')
 184                 chunks = var.split()
 185                 if len(chunks) > 1:
 186                     var = var.split()[0]
 187
 188                 # Environment vars the same as flag names without
 189                 # the initial -'s and in UPPERCASE.
 190                 env = var.strip('-').upper()
 191                 if env in os.environ:
 192                     if not is_flag_already_in_argv(var):
 193                         value = os.environ[env]
 194                         saved_messages.append(
 195                             f'Initialized from environment: {var} = {value}'
 196                         )
 197                         from string_utils import to_bool
 198                         if len(chunks) == 1 and to_bool(value):
 199                             sys.argv.append(var)
 200                         elif len(chunks) > 1:
 201                             sys.argv.append(var)
 202                             sys.argv.append(value)
 203                 var = ''
 204                 env = ''
 205         else:
 206             next
 207
 208     # Look for loadfile and read/parse it if present.
 209     loadfile = None
 210     saw_other_args = False
 211     grab_next_arg = False
 212     for arg in sys.argv[1:]:
 213         if 'config_loadfile' in arg:
 214             pieces = arg.split('=')
 215             if len(pieces) > 1:
 216                 loadfile = pieces[1]
 217             else:
 218                 grab_next_arg = True
 219         elif grab_next_arg:
 220             loadfile = arg
 221         else:
 222             saw_other_args = True
 223
 224     if loadfile is not None:
 225         if saw_other_args:
 226             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 227             print(msg, file=sys.stderr)
 228             saved_messages.append(msg)
 229         if not os.path.exists(loadfile):
 230             print(f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.',
 231                   file=sys.stderr)
 232             sys.exit(-1)
 233         with open(loadfile, 'r') as rf:
 234             newargs = rf.readlines()
 235         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 236         sys.argv += newargs
 237
 238     # Parse (possibly augmented, possibly completely overwritten)
 239     # commandline args with argparse normally and populate config.
 240     known, unknown = args.parse_known_args()
 241     config.update(vars(known))
 242
 243     # Reconstruct the argv with unrecognized flags for the benefit of
 244     # future argument parsers.  For example, unittest_main in python
 245     # has some of its own flags.  If we didn't recognize it, maybe
 246     # someone else will.
 247     sys.argv = sys.argv[:1] + unknown
 248
 249     # Check for savefile and populate it if requested.
 250     savefile = config['config_savefile']
 251     if savefile and len(savefile) > 0:
 252         with open(savefile, 'w') as wf:
 253             wf.write(
 254                 "\n".join(original_argv[1:])
 255             )
 256
 257     # Also dump the config on stderr if requested.
 258     if config['config_dump']:
 259         dump_config()
 260
 261     config_parse_called = True
 262     return config
 263
 264
 265 def has_been_parsed() -> bool:
 266     """Has the global config been parsed yet?"""
 267     global config_parse_called
 268     return config_parse_called
 269
 270
 271 def dump_config():
 272     """Print the current config to stdout."""
 273     print("Global Configuration:", file=sys.stderr)
 274     pprint.pprint(config, stream=sys.stderr)
 275     print()
 276
 277
 278 def late_logging():
 279     """Log messages saved earlier now that logging has been initialized."""
 280     logger = logging.getLogger(__name__)
 281     global saved_messages
 282     for _ in saved_messages:
 283         logger.debug(_)