edq.core.argparser

  1"""
  2A place to handle common CLI arguments.
  3"parsers" in this file are always assumed to be argparse parsers.
  4
  5The general idea is that callers can register callbacks to be called before and after parsing CLI arguments.
  6Pre-callbacks are generally intended to add arguments to the parser,
  7while post-callbacks are generally intended to act on the results of parsing.
  8"""
  9
 10import argparse
 11import functools
 12import typing
 13
 14import edq.core.config
 15import edq.core.log
 16import edq.util.net
 17
 18@typing.runtime_checkable
 19class PreParseFunction(typing.Protocol):
 20    """
 21    A function that can be called before parsing arguments.
 22    """
 23
 24    def __call__(self, parser: argparse.ArgumentParser, extra_state: typing.Dict[str, typing.Any]) -> None:
 25        """
 26        Prepare a parser for parsing.
 27        This is generally used for adding your module's arguments to the parser,
 28        for example a logging module may add arguments to set a logging level.
 29
 30        The extra state is shared between all pre-parse functions
 31        and will be placed in the final parsed output under `_pre_extra_state_`.
 32        """
 33
 34@typing.runtime_checkable
 35class PostParseFunction(typing.Protocol):
 36    """
 37    A function that can be called after parsing arguments.
 38    """
 39
 40    def __call__(self,
 41            parser: argparse.ArgumentParser,
 42            args: argparse.Namespace,
 43            extra_state: typing.Dict[str, typing.Any]) -> None:
 44        """
 45        Take actions after arguments are parsed.
 46        This is generally used for initializing your module with options,
 47        for example a logging module may set a logging level.
 48
 49        The extra state is shared between all post-parse functions
 50        and will be placed in the final parsed output under `_post_extra_state_`.
 51        """
 52
 53class Parser(argparse.ArgumentParser):
 54    """
 55    Extend an argparse parser to call the pre and post functions.
 56    """
 57
 58    def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None:
 59        super().__init__(*args, **kwargs)
 60
 61        self._pre_parse_callbacks: typing.Dict[str, PreParseFunction] = {}
 62        self._post_parse_callbacks: typing.Dict[str, PostParseFunction] = {}
 63
 64    def register_callbacks(self,
 65            key: str,
 66            pre_parse_callback: typing.Union[PreParseFunction, None] = None,
 67            post_parse_callback: typing.Union[PostParseFunction, None] = None,
 68            ) -> None:
 69        """
 70        Register callback functions to run before/after argument parsing.
 71        Any existing callbacks under the specified key will be replaced.
 72        """
 73
 74        if (pre_parse_callback is not None):
 75            self._pre_parse_callbacks[key] = pre_parse_callback
 76
 77        if (post_parse_callback is not None):
 78            self._post_parse_callbacks[key] = post_parse_callback
 79
 80    def parse_args(self,  # type: ignore[override]
 81            *args: typing.Any,
 82            skip_keys: typing.Union[typing.List[str], None] = None,
 83            **kwargs: typing.Any) -> argparse.Namespace:
 84        if (skip_keys is None):
 85            skip_keys = []
 86
 87        # Call pre-parse callbacks.
 88        pre_extra_state: typing.Dict[str, typing.Any] = {}
 89        for (key, pre_parse_callback) in self._pre_parse_callbacks.items():
 90            if (key not in skip_keys):
 91                pre_parse_callback(self, pre_extra_state)
 92
 93        # Parse the args.
 94        parsed_args = super().parse_args(*args, **kwargs)
 95
 96        # Call post-parse callbacks.
 97        post_extra_state: typing.Dict[str, typing.Any] = {}
 98        for (key, post_parse_callback) in self._post_parse_callbacks.items():
 99            if (key not in skip_keys):
100                post_parse_callback(self, parsed_args, post_extra_state)
101
102        # Attach the additional state to the args.
103        setattr(parsed_args, '_pre_extra_state_', pre_extra_state)
104        setattr(parsed_args, '_post_extra_state_', post_extra_state)
105
106        return parsed_args  # type: ignore[no-any-return]
107
108def get_default_parser(description: str,
109        version: typing.Union[str, None] = None,
110        include_log: bool = True,
111        include_config: bool = True,
112        include_net: bool = False,
113        config_options: typing.Union[typing.Dict[str, typing.Any], None] = None,
114        ) -> Parser:
115    """ Get a parser with the requested default callbacks already attached. """
116
117    if (config_options is None):
118        config_options = {}
119
120    parser = Parser(description = description)
121
122    if (version is not None):
123        parser.add_argument('--version',
124                action = 'version', version = version)
125
126    if (include_log):
127        parser.register_callbacks('log', edq.core.log.set_cli_args, edq.core.log.init_from_args)
128
129    if (include_config):
130        config_pre_func = functools.partial(edq.core.config.set_cli_args, **config_options)
131        config_post_func = functools.partial(edq.core.config.load_config_into_args, **config_options)
132        parser.register_callbacks('config', config_pre_func, config_post_func)
133
134    if (include_net):
135        parser.register_callbacks('net', edq.util.net.set_cli_args, edq.util.net.init_from_args)
136
137    return parser