config.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """Global configuration driven by commandline arguments, environment variables
   6 and saved configuration files.  This works across several modules.
   7
   8 Usage:
   9
  10     module.py:
  11     ----------
  12     import config
  13
  14     parser = config.add_commandline_args(
  15         "Module",
  16         "Args related to module doing the thing.",
  17     )
  18     parser.add_argument(
  19         "--module_do_the_thing",
  20         type=bool,
  21         default=True,
  22         help="Should the module do the thing?"
  23     )
  24
  25     main.py:
  26     --------
  27     import config
  28
  29     def main() -> None:
  30         parser = config.add_commandline_args(
  31             "Main",
  32             "A program that does the thing.",
  33         )
  34         parser.add_argument(
  35             "--dry_run",
  36             type=bool,
  37             default=False,
  38             help="Should we really do the thing?"
  39         )
  40         config.parse()   # Very important, this must be invoked!
  41
  42     If you set this up and remember to invoke config.parse(), all commandline
  43     arguments will play nicely together.  This is done automatically for you
  44     if you're using the bootstrap module's initialize wrapper.
  45
  46     % main.py -h
  47     usage: main.py [-h]
  48                    [--module_do_the_thing MODULE_DO_THE_THING]
  49                    [--dry_run DRY_RUN]
  50
  51     Module:
  52       Args related to module doing the thing.
  53
  54       --module_do_the_thing MODULE_DO_THE_THING
  55                    Should the module do the thing?
  56
  57     Main:
  58       A program that does the thing
  59
  60       --dry_run
  61                    Should we really do the thing?
  62
  63     Arguments themselves should be accessed via
  64     config.config['arg_name'].  e.g.
  65
  66     if not config.config['dry_run']:
  67         module.do_the_thing()
  68
  69 """
  70
  71 import argparse
  72 import logging
  73 import os
  74 import pprint
  75 import sys
  76 from typing import Any, Dict, List, Optional
  77
  78 # This module is commonly used by others in here and should avoid
  79 # taking any unnecessary dependencies back on them.
  80
  81 # Defer logging messages until later when logging has been initialized.
  82 SAVED_MESSAGES: List[str] = []
  83
  84 # Make a copy of the original program arguments.
  85 PROGRAM_NAME: str = os.path.basename(sys.argv[0])
  86 ORIG_ARGV: List[str] = sys.argv.copy()
  87
  88
  89 class OptionalRawFormatter(argparse.HelpFormatter):
  90     """This formatter has the same bahavior as the normal argparse text
  91     formatter except when the help text of an argument begins with
  92     "RAW|".  In that case, the line breaks are preserved and the text
  93     is not wrapped.
  94
  95     """
  96
  97     def _split_lines(self, text, width):
  98         if text.startswith('RAW|'):
  99             return text[4:].splitlines()
 100         return argparse.HelpFormatter._split_lines(self, text, width)
 101
 102
 103 # A global parser that we will collect arguments into.
 104 ARGS = argparse.ArgumentParser(
 105     description=None,
 106     formatter_class=OptionalRawFormatter,
 107     fromfile_prefix_chars="@",
 108     epilog=f'{PROGRAM_NAME} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.',
 109     # I don't fully understand why but when loaded by sphinx sometimes
 110     # the same module is loaded many times causing any arguments it
 111     # registers via module-level code to be redefined.  Work around
 112     # this iff the program is 'sphinx-build'
 113     conflict_handler='resolve' if PROGRAM_NAME == 'sphinx-build' else 'error',
 114 )
 115
 116 # Keep track of if we've been called and prevent being called more
 117 # than once.
 118 CONFIG_PARSE_CALLED = False
 119
 120
 121 # A global configuration dictionary that will contain parsed arguments.
 122 # It is also this variable that modules use to access parsed arguments.
 123 # This is the data that is most interesting to our callers; it will hold
 124 # the configuration result.
 125 config: Dict[str, Any] = {}
 126
 127 # It would be really nice if this shit worked from interactive python
 128
 129
 130 def add_commandline_args(title: str, description: str = ""):
 131     """Create a new context for arguments and return a handle."""
 132     return ARGS.add_argument_group(title, description)
 133
 134
 135 group = add_commandline_args(
 136     f'Global Config ({__file__})',
 137     'Args that control the global config itself; how meta!',
 138 )
 139 group.add_argument(
 140     '--config_loadfile',
 141     metavar='FILENAME',
 142     default=None,
 143     help='Config file (populated via --config_savefile) from which to read args in lieu or in addition to commandline.',
 144 )
 145 group.add_argument(
 146     '--config_dump',
 147     default=False,
 148     action='store_true',
 149     help='Display the global configuration (possibly derived from multiple sources) on STDERR at program startup.',
 150 )
 151 group.add_argument(
 152     '--config_savefile',
 153     type=str,
 154     metavar='FILENAME',
 155     default=None,
 156     help='Populate config file compatible with --config_loadfile to save global config for later use.',
 157 )
 158 group.add_argument(
 159     '--config_rejects_unrecognized_arguments',
 160     default=False,
 161     action='store_true',
 162     help=(
 163         'If present, config will raise an exception if it doesn\'t recognize an argument.  The '
 164         + 'default behavior is to ignore this so as to allow interoperability with programs that '
 165         + 'want to use their own argparse calls to parse their own, separate commandline args.'
 166     ),
 167 )
 168
 169
 170 def overwrite_argparse_epilog(msg: str) -> None:
 171     ARGS.epilog = msg
 172
 173
 174 def is_flag_already_in_argv(var: str):
 175     """Is a particular flag passed on the commandline?"""
 176     for _ in sys.argv:
 177         if var in _:
 178             return True
 179     return False
 180
 181
 182 def reorder_arg_action_groups_before_help(entry_module: Optional[str]):
 183     reordered_action_groups = []
 184     for grp in ARGS._action_groups:
 185         if entry_module is not None and entry_module in grp.title:  # type: ignore
 186             reordered_action_groups.append(grp)
 187         elif PROGRAM_NAME in group.title:  # type: ignore
 188             reordered_action_groups.append(grp)
 189         else:
 190             reordered_action_groups.insert(0, grp)
 191     return reordered_action_groups
 192
 193
 194 def augment_sys_argv_from_environment_variables():
 195     usage_message = ARGS.format_usage()
 196     optional = False
 197     var = ''
 198     for x in usage_message.split():
 199         if x[0] == '[':
 200             optional = True
 201         if optional:
 202             var += f'{x} '
 203             if x[-1] == ']':
 204                 optional = False
 205                 var = var.strip()
 206                 var = var.strip('[')
 207                 var = var.strip(']')
 208                 chunks = var.split()
 209                 if len(chunks) > 1:
 210                     var = var.split()[0]
 211
 212                 # Environment vars the same as flag names without
 213                 # the initial -'s and in UPPERCASE.
 214                 env = var.strip('-').upper()
 215                 if env in os.environ:
 216                     if not is_flag_already_in_argv(var):
 217                         value = os.environ[env]
 218                         SAVED_MESSAGES.append(f'Initialized from environment: {var} = {value}')
 219                         from string_utils import to_bool
 220
 221                         if len(chunks) == 1 and to_bool(value):
 222                             sys.argv.append(var)
 223                         elif len(chunks) > 1:
 224                             sys.argv.append(var)
 225                             sys.argv.append(value)
 226                 var = ''
 227                 env = ''
 228
 229
 230 def augment_sys_argv_from_loadfile():
 231     loadfile = None
 232     saw_other_args = False
 233     grab_next_arg = False
 234     for arg in sys.argv[1:]:
 235         if 'config_loadfile' in arg:
 236             pieces = arg.split('=')
 237             if len(pieces) > 1:
 238                 loadfile = pieces[1]
 239             else:
 240                 grab_next_arg = True
 241         elif grab_next_arg:
 242             loadfile = arg
 243         else:
 244             saw_other_args = True
 245
 246     if loadfile is not None:
 247         if not os.path.exists(loadfile):
 248             raise Exception(
 249                 f'ERROR: --config_loadfile argument must be a file, {loadfile} not found.'
 250             )
 251         if saw_other_args:
 252             msg = f'Augmenting commandline arguments with those from {loadfile}.'
 253         else:
 254             msg = f'Reading commandline arguments from {loadfile}.'
 255         print(msg, file=sys.stderr)
 256         SAVED_MESSAGES.append(msg)
 257
 258         with open(loadfile, 'r') as rf:
 259             newargs = rf.readlines()
 260         newargs = [arg.strip('\n') for arg in newargs if 'config_savefile' not in arg]
 261         sys.argv += newargs
 262
 263
 264 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 265     """Main program should call this early in main().  Note that the
 266     bootstrap.initialize wrapper takes care of this automatically.
 267
 268     """
 269     global CONFIG_PARSE_CALLED
 270     if CONFIG_PARSE_CALLED:
 271         return config
 272
 273     # If we're about to do the usage message dump, put the main
 274     # module's argument group last in the list (if possible) so that
 275     # when the user passes -h or --help, it will be visible on the
 276     # screen w/o scrolling.
 277     for arg in sys.argv:
 278         if arg in ('--help', '-h'):
 279             if entry_module is not None:
 280                 entry_module = os.path.basename(entry_module)
 281             ARGS._action_groups = reorder_arg_action_groups_before_help(entry_module)
 282
 283     # Examine the environment for variables that match known flags.
 284     # For a flag called --example_flag the corresponding environment
 285     # variable would be called EXAMPLE_FLAG.  If found, hackily add
 286     # these into sys.argv to be parsed.
 287     augment_sys_argv_from_environment_variables()
 288
 289     # Look for loadfile and read/parse it if present.  This also
 290     # works by jamming these values onto sys.argv.
 291     augment_sys_argv_from_loadfile()
 292
 293     # Parse (possibly augmented, possibly completely overwritten)
 294     # commandline args with argparse normally and populate config.
 295     known, unknown = ARGS.parse_known_args()
 296     config.update(vars(known))
 297
 298     # Reconstruct the argv with unrecognized flags for the benefit of
 299     # future argument parsers.  For example, unittest_main in python
 300     # has some of its own flags.  If we didn't recognize it, maybe
 301     # someone else will.
 302     if len(unknown) > 0:
 303         if config['config_rejects_unrecognized_arguments']:
 304             raise Exception(
 305                 f'Encountered unrecognized config argument(s) {unknown} with --config_rejects_unrecognized_arguments enabled; halting.'
 306             )
 307         SAVED_MESSAGES.append(f'Config encountered unrecognized commandline arguments: {unknown}')
 308     sys.argv = sys.argv[:1] + unknown
 309
 310     # Check for savefile and populate it if requested.
 311     savefile = config['config_savefile']
 312     if savefile and len(savefile) > 0:
 313         with open(savefile, 'w') as wf:
 314             wf.write("\n".join(ORIG_ARGV[1:]))
 315
 316     # Also dump the config on stderr if requested.
 317     if config['config_dump']:
 318         dump_config()
 319
 320     CONFIG_PARSE_CALLED = True
 321     return config
 322
 323
 324 def has_been_parsed() -> bool:
 325     """Has the global config been parsed yet?"""
 326     return CONFIG_PARSE_CALLED
 327
 328
 329 def dump_config():
 330     """Print the current config to stdout."""
 331     print("Global Configuration:", file=sys.stderr)
 332     pprint.pprint(config, stream=sys.stderr)
 333     print()
 334
 335
 336 def late_logging():
 337     """Log messages saved earlier now that logging has been initialized."""
 338     logger = logging.getLogger(__name__)
 339     logger.debug('Original commandline was: %s', ORIG_ARGV)
 340     for _ in SAVED_MESSAGES:
 341         logger.debug(_)