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