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