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