config.py

   1 #!/usr/bin/env python3
   2
   3 """Global configuration driven by commandline arguments (even across
   4 different modules).  Usage:
   5
   6     module.py:
   7     ----------
   8     import config
   9
  10     parser = config.add_commandline_args(
  11         "Module",
  12         "Args related to module doing the thing.",
  13     )
  14     parser.add_argument(
  15         "--module_do_the_thing",
  16         type=bool,
  17         default=True,
  18         help="Should the module do the thing?"
  19     )
  20
  21     main.py:
  22     --------
  23     import config
  24
  25     def main() -> None:
  26         parser = config.add_commandline_args(
  27             "Main",
  28             "A program that does the thing.",
  29         )
  30         parser.add_argument(
  31             "--dry_run",
  32             type=bool,
  33             default=False,
  34             help="Should we really do the thing?"
  35         )
  36         config.parse()   # Very important, this must be invoked!
  37
  38     If you set this up and remember to invoke config.parse(), all commandline
  39     arguments will play nicely together:
  40
  41     % main.py -h
  42     usage: main.py [-h]
  43                    [--module_do_the_thing MODULE_DO_THE_THING]
  44                    [--dry_run DRY_RUN]
  45
  46     Module:
  47       Args related to module doing the thing.
  48
  49       --module_do_the_thing MODULE_DO_THE_THING
  50                    Should the module do the thing?
  51
  52     Main:
  53       A program that does the thing
  54
  55       --dry_run
  56                    Should we really do the thing?
  57
  58     Arguments themselves should be accessed via
  59     config.config['arg_name'].  e.g.
  60
  61     if not config.config['dry_run']:
  62         module.do_the_thing()
  63
  64 """
  65
  66 import argparse
  67 import logging
  68 import os
  69 import pprint
  70 import sys
  71 from typing import Any, Dict, List, Optional
  72
  73 # This module is commonly used by others in here and should avoid
  74 # taking any unnecessary dependencies back on them.
  75
  76 # Note: at this point in time, logging hasn't been configured and
  77 # anything we log will come out the root logger.
  78 logger = logging.getLogger(__name__)
  79
  80 # Defer logging messages until later when logging has been initialized.
  81 saved_messages: List[str] = []
  82
  83 # Make a copy of the original program arguments.
  84 program_name = os.path.basename(sys.argv[0])
  85 original_argv = [arg for arg in sys.argv]
  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 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 --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         logger.warning('config.parse has already been called; ignoring spurious invocation')
 149         return config
 150
 151     global saved_messages
 152
 153     # If we're about to do the usage message dump, put the main module's
 154     # argument group last in the list (if possible) so that when the user
 155     # passes -h or --help, it will be visible on the screen w/o scrolling.
 156     reordered_action_groups = []
 157     global prog
 158     for arg in sys.argv:
 159         if arg == '--help' or arg == '-h':
 160             for group in args._action_groups:
 161                 if entry_module is not None and entry_module in group.title:
 162                     reordered_action_groups.append(group)
 163                 elif program_name in group.title:
 164                     reordered_action_groups.append(group)
 165                 else:
 166                     reordered_action_groups.insert(0, group)
 167             args._action_groups = reordered_action_groups
 168
 169     # Examine the environment variables that match known flags.  For a
 170     # flag called --example_flag the corresponding environment
 171     # variable would be called EXAMPLE_FLAG.
 172     usage_message = args.format_usage()
 173     optional = False
 174     var = ''
 175     for x in usage_message.split():
 176         if x[0] == '[':
 177             optional = True
 178         if optional:
 179             var += f'{x} '
 180             if x[-1] == ']':
 181                 optional = False
 182                 var = var.strip()
 183                 var = var.strip('[')
 184                 var = var.strip(']')
 185                 chunks = var.split()
 186                 if len(chunks) > 1:
 187                     var = var.split()[0]
 188
 189                 # Environment vars the same as flag names without
 190                 # the initial -'s and in UPPERCASE.
 191                 env = var.strip('-').upper()
 192                 if env in os.environ:
 193                     if not is_flag_already_in_argv(var):
 194                         value = os.environ[env]
 195                         saved_messages.append(
 196                             f'Initialized from environment: {var} = {value}'
 197                         )
 198                         from string_utils import to_bool
 199                         if len(chunks) == 1 and to_bool(value):
 200                             sys.argv.append(var)
 201                         elif len(chunks) > 1:
 202                             sys.argv.append(var)
 203                             sys.argv.append(value)
 204                 var = ''
 205                 env = ''
 206         else:
 207             next
 208
 209     # Look for loadfile and read/parse it if present.
 210     loadfile = None
 211     saw_other_args = False
 212     grab_next_arg = False
 213     for arg in sys.argv[1:]:
 214         if 'config_loadfile' in arg:
 215             pieces = arg.split('=')
 216             if len(pieces) > 1:
 217                 loadfile = pieces[1]
 218             else:
 219                 grab_next_arg = True
 220         elif grab_next_arg:
 221             loadfile = arg
 222         else:
 223             saw_other_args = True
 224
 225     if loadfile is not None:
 226         if saw_other_args:
 227             msg = f'WARNING: Augmenting commandline arguments with those from {loadfile}.'
 228             print(msg, file=sys.stderr)
 229             saved_messages.append(msg)
 230         if not os.path.exists(loadfile):
 231             print(f'--config_loadfile argument must be a file, {loadfile} not found.',
 232                   file=sys.stderr)
 233             sys.exit(-1)
 234         with open(loadfile, 'r') as rf:
 235             newargs = rf.readlines()
 236         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 237         sys.argv += newargs
 238
 239     # Parse (possibly augmented, possibly completely overwritten)
 240     # commandline args with argparse normally and populate config.
 241     known, unknown = args.parse_known_args()
 242     config.update(vars(known))
 243
 244     # Reconstruct the argv with unrecognized flags for the benefit of
 245     # future argument parsers.  For example, unittest_main in python
 246     # has some of its own flags.  If we didn't recognize it, maybe
 247     # someone else will.
 248     sys.argv = sys.argv[:1] + unknown
 249
 250     # Check for savefile and populate it if requested.
 251     savefile = config['config_savefile']
 252     if savefile and len(savefile) > 0:
 253         with open(savefile, 'w') as wf:
 254             wf.write(
 255                 "\n".join(original_argv[1:])
 256             )
 257
 258     # Also dump the config on stderr if requested.
 259     if config['config_dump']:
 260         dump_config()
 261
 262     config_parse_called = True
 263     return config
 264
 265
 266 def has_been_parsed() -> bool:
 267     """Has the global config been parsed yet?"""
 268     global config_parse_called
 269     return config_parse_called
 270
 271
 272 def dump_config():
 273     """Print the current config to stdout."""
 274     print("Global Configuration:", file=sys.stderr)
 275     pprint.pprint(config, stream=sys.stderr)
 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(_)