config.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312
  1. #
  2. # Copyright (c) 2021 Project CHIP Authors
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. #
  16. """Configuration utilities for MDF tools"""
  17. import argparse
  18. import ast
  19. import logging
  20. import re
  21. from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Pattern, Sequence, Tuple, Union
  22. import humanfriendly # type: ignore
  23. import memdf.util.nd as nd
  24. import memdf.util.pretty
  25. # A ConfigDescription is a declarative description of configuration options.
  26. #
  27. # In a description dictionary, (most) keys are configuration keys
  28. # and values are dictionaries that MUST contain at least
  29. # 'help': help string.
  30. # 'default': default value.
  31. # and may contain:
  32. # 'metavar': if the command line argument takes a value
  33. # 'choices': if the argument value must be one of several specific values
  34. # 'argparse': additional argument parsing information; most of this is
  35. # supplied as keyword arguments to `argparse.add_argument()`,
  36. # except for:
  37. # 'alias': list of alternate command line option names
  38. # 'postprocess': a callable invoked after argument parsing with three
  39. # arguments: the config, the key, and the description entry.
  40. #
  41. # Special keys can be used to control argument parser groups. By default any
  42. # configuration key containing a ‘.’ belongs to a group determined by the
  43. # key prefix (the part before the first ‘.’).
  44. # Config.group_def(group):
  45. # the value is supplied as keyword arguments to
  46. # `argparse.add_argument_group()`
  47. # Config.group_map(prefix):
  48. # the value contains a key 'group', whose value is the group
  49. # to be used for configuration keys with the given prefix.
  50. #
  51. ConfigDescription = Mapping[Union[str, Tuple[int, str]], Mapping[str, Any]]
  52. class Config:
  53. """Key/value store and argument parsing.
  54. A configuration key is a string where dots (`.`) separate levels in the
  55. underlying nested dictionary.
  56. For functions that take a Config, an empty `Config()` is normally
  57. acceptable. These functions should always assume reasonable defaults,
  58. so that they can be used without any particular configuration.
  59. """
  60. def __init__(self):
  61. self.d: MutableMapping = {}
  62. self.argparse = None
  63. self.argparse_groups = {}
  64. self.group_alias = {}
  65. self.postprocess_args = {}
  66. self.config_desc: ConfigDescription = None
  67. self.dest_to_key: MutableMapping = {}
  68. self.key_to_dest: MutableMapping = {}
  69. # Basic config access
  70. def get(self, key: str, default: Any = None) -> Any:
  71. return self.getl(key.split('.'), default)
  72. def __getitem__(self, key: str) -> Any:
  73. """[] syntax for configuration.
  74. Note that this will return `None` for an unknown key, since the
  75. absence of a configured value is not considered an error.
  76. """
  77. return self.get(key)
  78. def getl(self, keys: nd.Key, default: Any = None) -> Any:
  79. return nd.get(self.d, keys, default)
  80. def put(self, key: str, value: Any) -> None:
  81. self.putl(key.split('.'), value)
  82. def __setitem__(self, key: str, value: Any) -> None:
  83. self.put(key, value)
  84. def putl(self, keys: nd.Key, value: Any) -> None:
  85. nd.put(self.d, keys, value)
  86. def update(self, src: Mapping) -> None:
  87. nd.update(self.d, src)
  88. # Command line and config file reading
  89. _GROUP_DEF = 1
  90. _GROUP_MAP = 2
  91. @staticmethod
  92. def group_def(s: str) -> Tuple[int, str]:
  93. return (Config._GROUP_DEF, s)
  94. @staticmethod
  95. def group_map(s: str) -> Tuple[int, str]:
  96. return (Config._GROUP_MAP, s)
  97. def init_config(self, desc: ConfigDescription) -> 'Config':
  98. """Initialize a configuration from a description dictionary.
  99. Note that this initializes only the key/value store,
  100. not anything associated with command line argument parsing.
  101. """
  102. self.config_desc = desc
  103. for key, info in desc.items():
  104. if isinstance(key, str):
  105. self.put(key, info['default'])
  106. return self
  107. def init_args(self, desc: ConfigDescription, *args, **kwargs) -> 'Config':
  108. """Initialize command line argument parsing."""
  109. self.argparse = argparse.ArgumentParser(*args, **kwargs)
  110. # Groups
  111. for key, info in desc.items():
  112. if not isinstance(key, tuple):
  113. continue
  114. kind, name = key
  115. if kind == self._GROUP_MAP:
  116. self.group_alias[name] = info['group']
  117. elif kind == self._GROUP_DEF:
  118. self.argparse_groups[name] = self.argparse.add_argument_group(
  119. **info)
  120. # Arguments
  121. for key, info in desc.items():
  122. if not isinstance(key, str):
  123. continue
  124. if (arg_info := info.get('argparse', {})) is False:
  125. continue
  126. arg_info = arg_info.copy()
  127. name = arg_info.pop('argument', '--' + key.replace('.', '-'))
  128. names = [name] + arg_info.pop('alias', [])
  129. info['names'] = names
  130. for k in ['metavar', 'choices']:
  131. if k in info:
  132. arg_info[k] = info[k]
  133. default = info['default']
  134. if not arg_info.get('action'):
  135. if isinstance(default, list):
  136. arg_info['action'] = 'append'
  137. elif default is False:
  138. arg_info['action'] = 'store_true'
  139. elif default is True:
  140. arg_info['action'] = 'store_false'
  141. elif isinstance(default, int) and 'metavar' not in info:
  142. arg_info['action'] = 'count'
  143. if postprocess := info.get('postprocess'):
  144. self.postprocess_args[key] = (postprocess, info)
  145. group: Optional[str] = info.get('group')
  146. if group is None and (e := key.find('.')) > 0:
  147. group = key[0:e]
  148. group = self.group_alias.get(group, group)
  149. arg_group = self.argparse_groups.get(group, self.argparse)
  150. arg = arg_group.add_argument(*names,
  151. help=info['help'],
  152. default=self.get(key, default),
  153. **arg_info)
  154. self.dest_to_key[arg.dest] = key
  155. self.key_to_dest[key] = arg.dest
  156. return self
  157. def init(self, desc: ConfigDescription, *args, **kwargs) -> 'Config':
  158. """Intialize configuration from a configuration description."""
  159. self.init_config(desc)
  160. self.init_args(desc, *args, **kwargs)
  161. return self
  162. def parse(self, argv: Sequence[str]) -> 'Config':
  163. """Parse command line options into a configuration dictionary."""
  164. # Read config file(s).
  165. config_parser = argparse.ArgumentParser(add_help=False,
  166. allow_abbrev=False)
  167. config_arg: Dict[str, Any] = {
  168. 'metavar': 'FILE',
  169. 'default': [],
  170. 'action': 'append',
  171. 'help': 'Read configuration FILE'
  172. }
  173. config_parser.add_argument('--config-file', **config_arg)
  174. self.argparse.add_argument('--config-file', **config_arg)
  175. config_args, argv = config_parser.parse_known_args(argv[1:])
  176. for filename in config_args.config_file:
  177. self.read_config_file(filename)
  178. # Update argparser defaults.
  179. defaults = {}
  180. for dest, key in self.dest_to_key.items():
  181. default = self.get(key)
  182. if default is not None:
  183. defaults[dest] = default
  184. self.argparse.set_defaults(**defaults)
  185. # Parse command line arguments and update config.
  186. args = self.argparse.parse_args(argv)
  187. for dest, value in vars(args).items():
  188. if (key := self.dest_to_key.get(dest)) is None:
  189. key = 'args.' + dest
  190. self.put(key, value)
  191. # Configure logging.
  192. if self.get('log-level') is None:
  193. verbose = self.get('verbose', 0)
  194. self.put('log-level',
  195. (logging.DEBUG if verbose > 1 else
  196. logging.INFO if verbose else logging.WARNING))
  197. else:
  198. self.put('log-level',
  199. getattr(logging, self.get('log-level').upper()))
  200. logging.basicConfig(level=self.get('log-level'),
  201. format=self.get('log-format'))
  202. # Postprocess config.
  203. for key, postprocess in self.postprocess_args.items():
  204. action, info = postprocess
  205. action(self, key, info)
  206. memdf.util.pretty.debug(self.d)
  207. return self
  208. def read_config_file(self, filename: str) -> 'Config':
  209. """Read a configuration file."""
  210. with open(filename, 'r') as fp:
  211. d = ast.literal_eval(fp.read())
  212. nd.update(self.d, d)
  213. return self
  214. @staticmethod
  215. def transpose_dictlist(src: Dict[str, List[str]]) -> Dict[str, str]:
  216. d: Dict[str, str] = {}
  217. for k, vlist in src.items():
  218. for v in vlist:
  219. d[v] = k
  220. return d
  221. def getl_re(self, key: nd.Key) -> Optional[Pattern]:
  222. """Get a cached compiled regular expression for a config value list."""
  223. regex_key: nd.Key = ['cache', 're'] + key
  224. regex: Optional[Pattern] = self.getl(regex_key)
  225. if not regex:
  226. branches: Optional[Sequence[str]] = self.getl(key)
  227. if branches:
  228. regex = re.compile('|'.join(branches))
  229. self.putl(regex_key, regex)
  230. return regex
  231. def get_re(self, key: str) -> Optional[Pattern]:
  232. return self.getl_re(key.split('.'))
  233. # Argument parsing helpers
  234. def parse_size(s: str) -> int:
  235. return humanfriendly.parse_size(s, binary=True) if s else 0
  236. class ParseSizeAction(argparse.Action):
  237. """argparse helper for humanfriendly sizes"""
  238. def __call__(self, parser, namespace, values, option_string=None):
  239. setattr(namespace, self.dest, parse_size(values))
  240. # Config description of options shared by all tools.
  241. CONFIG: ConfigDescription = {
  242. 'log-level': {
  243. 'help':
  244. 'Set logging level: one of critical, error, warning, info, debug.',
  245. 'default': None,
  246. 'choices': ['critical', 'error', 'warning', 'info', 'debug'],
  247. },
  248. 'log-format': {
  249. 'help': 'Set logging format',
  250. 'metavar': 'FORMAT',
  251. 'default': '%(message)s',
  252. },
  253. 'verbose': {
  254. 'help': 'Show informational messages; repeat for debugging messages',
  255. 'default': 0,
  256. 'argparse': {
  257. 'alias': ['-v'],
  258. 'action': 'count',
  259. },
  260. },
  261. }