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