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
  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: 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, please.
 170     reordered_action_groups = []
 171     prog = sys.argv[0]
 172     for arg in sys.argv:
 173         if arg == '--help' or arg == '-h':
 174             print(entry_module)
 175             for group in args._action_groups:
 176                 if entry_module in group.title or prog in group.title:
 177                     reordered_action_groups.insert(0, group)
 178                 else:
 179                     reordered_action_groups.append(group)
 180             args._action_groups = reordered_action_groups
 181
 182     # Examine the environment variables to settings that match
 183     # known flags.  For a flag called --example_flag the corresponding
 184     # environment variable would be called EXAMPLE_FLAG.
 185     usage_message = args.format_usage()
 186     optional = False
 187     var = ''
 188     for x in usage_message.split():
 189         if x[0] == '[':
 190             optional = True
 191         if optional:
 192             var += f'{x} '
 193             if x[-1] == ']':
 194                 optional = False
 195                 var = var.strip()
 196                 var = var.strip('[')
 197                 var = var.strip(']')
 198                 chunks = var.split()
 199                 if len(chunks) > 1:
 200                     var = var.split()[0]
 201
 202                 # Environment vars the same as flag names without
 203                 # the initial -'s and in UPPERCASE.
 204                 env = var.strip('-').upper()
 205                 if env in os.environ:
 206                     if not is_flag_already_in_argv(var):
 207                         value = os.environ[env]
 208                         saved_messages.append(
 209                             f'Initialized from environment: {var} = {value}'
 210                         )
 211                         from string_utils import to_bool
 212                         if len(chunks) == 1 and to_bool(value):
 213                             sys.argv.append(var)
 214                         elif len(chunks) > 1:
 215                             sys.argv.append(var)
 216                             sys.argv.append(value)
 217                 var = ''
 218                 env = ''
 219         else:
 220             next
 221
 222     # Parse (possibly augmented) commandline args with argparse normally.
 223     known, unknown = args.parse_known_args()
 224     config.update(vars(known))
 225
 226     # Reconstruct the argv with unrecognized flags for the benefit of
 227     # future argument parsers.  For example, unittest_main in python
 228     # has some of its own flags.  If we didn't recognize it, maybe
 229     # someone else will.
 230     sys.argv = sys.argv[:1] + unknown
 231
 232     if config['config_savefile']:
 233         with open(config['config_savefile'], 'w') as wf:
 234             wf.write("\n".join(sys.argv[1:]))
 235
 236     if config['config_dump']:
 237         dump_config()
 238     return config
 239
 240
 241 def has_been_parsed() -> bool:
 242     """Has the global config been parsed yet?"""
 243     global config_parse_called
 244     return config_parse_called
 245
 246
 247 def dump_config():
 248     """Print the current config to stdout."""
 249     print("Global Configuration:", file=sys.stderr)
 250     pprint.pprint(config, stream=sys.stderr)
 251
 252
 253 def late_logging():
 254     """Log messages saved earlier now that logging has been initialized."""
 255     logger = logging.getLogger(__name__)
 256     global saved_messages
 257     for _ in saved_messages:
 258         logger.debug(_)