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 config.config['arg_name'].  e.g.
  59
  60     if not config.config['dry_run']:
  61         module.do_the_thing()
  62 """
  63
  64 import argparse
  65 import pprint
  66 import re
  67 import sys
  68 from typing import Dict, Any
  69
  70 # Note: at this point in time, logging hasn't been configured and
  71 # anything we log will come out the root logger.
  72
  73
  74 class LoadFromFile(argparse.Action):
  75     """Helper to load a config file into argparse."""
  76     def __call__ (self, parser, namespace, values, option_string = None):
  77         with values as f:
  78             buf = f.read()
  79             argv = []
  80             for line in buf.split(','):
  81                 line = line.strip()
  82                 line = line.strip('{')
  83                 line = line.strip('}')
  84                 m = re.match(r"^'([a-zA-Z_\-]+)'\s*:\s*(.*)$", line)
  85                 if m:
  86                     key = m.group(1)
  87                     value = m.group(2)
  88                     value = value.strip("'")
  89                     if value not in ('None', 'True', 'False'):
  90                         argv.append(f'--{key}')
  91                         argv.append(value)
  92             parser.parse_args(argv, namespace)
  93
  94
  95 # A global parser that we will collect arguments into.
  96 args = argparse.ArgumentParser(
  97     description=f"This program uses config.py ({__file__}) for global, cross-module configuration.",
  98     formatter_class=argparse.ArgumentDefaultsHelpFormatter,
  99     fromfile_prefix_chars="@"
 100 )
 101 config_parse_called = False
 102
 103 # A global configuration dictionary that will contain parsed arguments
 104 # It is also this variable that modules use to access parsed arguments
 105 config: Dict[str, Any] = {}
 106
 107
 108 def add_commandline_args(title: str, description: str = ""):
 109     """Create a new context for arguments and return a handle."""
 110     return args.add_argument_group(title, description)
 111
 112
 113 group = add_commandline_args(
 114     f'Global Config ({__file__})',
 115     'Args that control the global config itself; how meta!',
 116 )
 117 group.add_argument(
 118     '--config_loadfile',
 119     type=open,
 120     action=LoadFromFile,
 121     metavar='FILENAME',
 122     default=None,
 123     help='Config file from which to read args in lieu or in addition to commandline.',
 124 )
 125 group.add_argument(
 126     '--config_dump',
 127     default=False,
 128     action='store_true',
 129     help='Display the global configuration on STDERR at program startup.',
 130 )
 131 group.add_argument(
 132     '--config_savefile',
 133     type=str,
 134     metavar='FILENAME',
 135     default=None,
 136     help='Populate config file compatible --config_loadfile to save config for later use.',
 137 )
 138
 139
 140 def parse() -> Dict[str, Any]:
 141     """Main program should call this early in main()"""
 142     global config_parse_called
 143     config_parse_called = True
 144     config.update(vars(args.parse_args()))
 145
 146     if config['config_savefile']:
 147         with open(config['config_savefile'], 'w') as wf:
 148             wf.write("\n".join(sys.argv[1:]))
 149
 150     if config['config_dump']:
 151         dump_config()
 152
 153     return config
 154
 155
 156 def has_been_parsed() -> bool:
 157     global config_parse_called
 158     return config_parse_called
 159
 160
 161 def dump_config():
 162     """Print the current config to stdout."""
 163     print("Global Configuration:", file=sys.stderr)
 164     pprint.pprint(config, stream=sys.stderr)