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 re
  71 import sys
  72 from typing import Any, Dict, List, Optional
  73
  74 # This module is commonly used by others in here and should avoid
  75 # taking any unnecessary dependencies back on them.
  76
  77 # Note: at this point in time, logging hasn't been configured and
  78 # anything we log will come out the root logger.
  79
  80
  81 class LoadFromFile(argparse.Action):
  82     """Helper to load a config file into argparse."""
  83     def __call__ (self, parser, namespace, values, option_string = None):
  84         with values as f:
  85             buf = f.read()
  86             argv = []
  87             for line in buf.split(','):
  88                 line = line.strip()
  89                 line = line.strip('{')
  90                 line = line.strip('}')
  91                 m = re.match(r"^'([a-zA-Z_\-]+)'\s*:\s*(.*)$", line)
  92                 if m:
  93                     key = m.group(1)
  94                     value = m.group(2)
  95                     value = value.strip("'")
  96                     if value not in ('None', 'True', 'False'):
  97                         argv.append(f'--{key}')
  98                         argv.append(value)
  99             parser.parse_args(argv, namespace)
 100
 101
 102 # A global parser that we will collect arguments into.
 103 prog = os.path.basename(sys.argv[0])
 104 args = argparse.ArgumentParser(
 105     description=None,
 106     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
 107     fromfile_prefix_chars="@",
 108     epilog=f'-----------------------------------------------------------------------------\n{prog} uses config.py ({__file__}) for global, cross-module configuration setup and parsing.\n-----------------------------------------------------------------------------'
 109 )
 110 config_parse_called = False
 111
 112 # A global configuration dictionary that will contain parsed arguments
 113 # It is also this variable that modules use to access parsed arguments
 114 config: Dict[str, Any] = {}
 115
 116 # Defer logging messages until later when logging has been initialized.
 117 saved_messages: List[str] = []
 118
 119
 120 def add_commandline_args(title: str, description: str = ""):
 121     """Create a new context for arguments and return a handle."""
 122     return args.add_argument_group(title, description)
 123
 124
 125 group = add_commandline_args(
 126     f'Global Config ({__file__})',
 127     'Args that control the global config itself; how meta!',
 128 )
 129 group.add_argument(
 130     '--config_loadfile',
 131     type=open,
 132     action=LoadFromFile,
 133     metavar='FILENAME',
 134     default=None,
 135     help='Config file from which to read args in lieu or in addition to commandline.',
 136 )
 137 group.add_argument(
 138     '--config_dump',
 139     default=False,
 140     action='store_true',
 141     help='Display the global configuration on STDERR at program startup.',
 142 )
 143 group.add_argument(
 144     '--config_savefile',
 145     type=str,
 146     metavar='FILENAME',
 147     default=None,
 148     help='Populate config file compatible --config_loadfile to save config for later use.',
 149 )
 150
 151
 152 def is_flag_already_in_argv(var: str):
 153     """Is a particular flag passed on the commandline?"""
 154     for _ in sys.argv:
 155         if var in _:
 156             return True
 157     return False
 158
 159
 160 def parse(entry_module: Optional[str]) -> Dict[str, Any]:
 161     """Main program should call this early in main()"""
 162     global config_parse_called
 163     if config_parse_called:
 164         return config
 165     config_parse_called = True
 166     global saved_messages
 167
 168     # If we're about to do the usage message dump, put the main module's
 169     # argument group first in the list (if possible), please.
 170     reordered_action_groups = []
 171     prog = sys.argv[0]
 172
 173     for arg in sys.argv:
 174         if arg == '--help' or arg == '-h':
 175             for group in args._action_groups:
 176                 if entry_module is not None and entry_module in group.title:
 177                     reordered_action_groups.insert(0, group)  # prepend
 178                 elif prog in group.title:
 179                     reordered_action_groups.insert(0, group)  # prepend
 180                 else:
 181                     reordered_action_groups.append(group)     # append
 182             args._action_groups = reordered_action_groups
 183
 184     # Examine the environment variables to settings that match
 185     # known flags.  For a flag called --example_flag the corresponding
 186     # environment variable would be called EXAMPLE_FLAG.
 187     usage_message = args.format_usage()
 188     optional = False
 189     var = ''
 190     for x in usage_message.split():
 191         if x[0] == '[':
 192             optional = True
 193         if optional:
 194             var += f'{x} '
 195             if x[-1] == ']':
 196                 optional = False
 197                 var = var.strip()
 198                 var = var.strip('[')
 199                 var = var.strip(']')
 200                 chunks = var.split()
 201                 if len(chunks) > 1:
 202                     var = var.split()[0]
 203
 204                 # Environment vars the same as flag names without
 205                 # the initial -'s and in UPPERCASE.
 206                 env = var.strip('-').upper()
 207                 if env in os.environ:
 208                     if not is_flag_already_in_argv(var):
 209                         value = os.environ[env]
 210                         saved_messages.append(
 211                             f'Initialized from environment: {var} = {value}'
 212                         )
 213                         from string_utils import to_bool
 214                         if len(chunks) == 1 and to_bool(value):
 215                             sys.argv.append(var)
 216                         elif len(chunks) > 1:
 217                             sys.argv.append(var)
 218                             sys.argv.append(value)
 219                 var = ''
 220                 env = ''
 221         else:
 222             next
 223
 224     # Parse (possibly augmented) commandline args with argparse normally.
 225     known, unknown = args.parse_known_args()
 226     config.update(vars(known))
 227
 228     # Reconstruct the argv with unrecognized flags for the benefit of
 229     # future argument parsers.  For example, unittest_main in python
 230     # has some of its own flags.  If we didn't recognize it, maybe
 231     # someone else will.
 232     sys.argv = sys.argv[:1] + unknown
 233
 234     if config['config_savefile']:
 235         with open(config['config_savefile'], 'w') as wf:
 236             wf.write("\n".join(sys.argv[1:]))
 237
 238     if config['config_dump']:
 239         dump_config()
 240     return config
 241
 242
 243 def has_been_parsed() -> bool:
 244     """Has the global config been parsed yet?"""
 245     global config_parse_called
 246     return config_parse_called
 247
 248
 249 def dump_config():
 250     """Print the current config to stdout."""
 251     print("Global Configuration:", file=sys.stderr)
 252     pprint.pprint(config, stream=sys.stderr)
 253
 254
 255 def late_logging():
 256     """Log messages saved earlier now that logging has been initialized."""
 257     logger = logging.getLogger(__name__)
 258     global saved_messages
 259     for _ in saved_messages:
 260         logger.debug(_)