Source code for argus_cli.arguments

import collections
import inspect
import re
import sys

import argparse
from argparse import ArgumentParser
from pydoc import locate

from argus_cli.helpers.collections import ImmutableDeepDict
from argus_cli.helpers.log import log
from argus_cli.helpers.formatting import to_caterpillar_case

#: The key of the plugin argument
_PLUGIN_ARGUMENT = "_plugins"
# Storage location for the parsers
_PLUGIN_PARSER = "_parser"
_PLUGIN_SUBPARSER = "_subparser"

# Regexes used to check for arguments
#: Checks if the line is argument metadata
_ARGUMENT_META_REGEX = re.compile(r":(?:param|alias)")
#: Gets data from the parameter
_PARAM_REGEX = re.compile(r":param (?P<name>[*\w]+): (?P<doc>.*)")
#: Gets aliases for parameter
_ALIAS_REGEX = re.compile(r":alias\s+(?P<name>[*\w]+):\s+(?P<aliases>.*)")
#: Gets data from the parameter
_TYPE_PARAM_REGEX = re.compile(r":param (?P<argument_type>[*\w]+) (?P<name>[*\w]+): (?P<doc>.*)")


[docs]class PluginParserContainer(object): """A container that handles plugin parsers Each node in the dict has three objects. _subparsers: The _SubParserAction object (that you add parsers to) _parser: The actual parser. This is just used for testing Everything else: Commands and sub-plugins """ class _StorePlugin(argparse._SubParsersAction): """Makes plugin subparsers store the used subparser in one list.""" def __call__(self, parser, namespace, values, option_string=None): super().__call__(parser, namespace, values, option_string) # Values will be all arguments after this plugin. # Just get the name of this plugin. values = values[0] if not hasattr(namespace, _PLUGIN_ARGUMENT): setattr(namespace, _PLUGIN_ARGUMENT, [values]) else: getattr(namespace, _PLUGIN_ARGUMENT).insert(0, values) def __init__(self, main_parser): self._dict = ImmutableDeepDict() self._plugin_parser = main_parser.add_subparsers( action=self._StorePlugin, help="Which plugin to use", metavar="<Plugin [sub-plugin ...]>" ) def __str__(self): return str(self._dict) def _add_parser(self, plugin: tuple) -> ArgumentParser: """Adds a plugin's parser to the tree.""" plugin_name = plugin[-1] if len(plugin) == 1: # This is a top level plugin parent_parser = self._plugin_parser else: try: parent_parser = self._dict[plugin[:-1]][_PLUGIN_SUBPARSER] except KeyError: # A key-error means that the parent doesn't have a parser. parent_parser = self._add_parser(plugin[:-1]) new_parser = parent_parser.add_parser(plugin_name) new_subparser = new_parser.add_subparsers(action=self._StorePlugin) self._dict[plugin] = { _PLUGIN_PARSER: new_parser, # Used for testing _PLUGIN_SUBPARSER: new_subparser } return new_subparser
[docs] def add_parser(self, plugin: tuple) -> None: """Adds a parser to the container""" self._add_parser(plugin)
[docs] def add_command(self, plugin: tuple, command_name: str) -> ArgumentParser: """Adds a command to a plugin parser""" plugin_container = self._dict[plugin] plugin_parser = plugin_container[_PLUGIN_SUBPARSER] plugin_container[command_name] = plugin_parser.add_parser(command_name) return plugin_container[command_name]
[docs] def get_command(self, plugin: tuple, command_name: str): """Gets a command from a plugin""" return self._dict[plugin][command_name]
#: The parser to rule them all! _ROOT_PARSER = ArgumentParser(prog="Argus Toolbelt") #: Parsers that handles a plugins commands _PARSERS = PluginParserContainer(_ROOT_PARSER) def _parse_type(parameter_type, argument: dict) -> dict: """Parses a parameter's type :param parameter_type: The type of the parameter :param argument: The argument to modify :returns: Modified argument """ if isinstance(parameter_type, list): if any(isinstance(element, collections.Sequence) and not isinstance(element, str) for element in parameter_type): raise ValueError("A list of choices can not have a nested iterable object.") argument["choices"] = parameter_type elif inspect.isclass(parameter_type) and issubclass(parameter_type, (list, tuple)): # When a user specifies a "list" input it's not a instance of list yet. # Hence here we use issubclass() instead. argument["nargs"] = '*' elif isinstance(parameter_type, collections.Callable): argument["type"] = parameter_type else: raise ValueError("Non-supported parameter type") return argument def _parse_parameters(function: callable) -> dict: """Parses a functions parameters. :param function: The function to parse :returns: All arguments, ready to pass to argparse """ arguments = collections.OrderedDict() signature = inspect.signature(function) log.debug("%s arguments: %s" % (function.__name__, signature)) for name, parameter in signature.parameters.items(): if parameter.kind == inspect.Parameter.VAR_KEYWORD: # **kwargs Have to be defined in docstrings log.debug("**%s argument ignored. kwargs are added from the docstring." % name) continue arguments[name] = {"names": [to_caterpillar_case(name)]} if parameter.default is not parameter.empty: arguments[name]["required"] = False arguments[name]["default"] = parameter.default if parameter.kind == inspect.Parameter.VAR_POSITIONAL: arguments[name]["required"] = False arguments[name]["nargs"] = "*" if parameter.annotation is not parameter.empty: arguments[name] = _parse_type(parameter.annotation, arguments[name]) log.debug("%s: Registered commands from signature:\n\t%s" % (function.__name__, arguments)) return arguments def _parse_docstring(function: callable, parsed: dict) -> dict: """Parses a function's docstring for more info about it and it's parameters. :param function: The function to parse :param parsed: Existing arguments for the function :return: Short description, Long description and more argument info """ # Escape all % { and } so argparse doesnt crash when trying to format the string lines = function.__doc__.replace("{", "{{").replace("}", "}}").replace("%", "%%").split("\n", 1) help_text = lines[0].strip() if len(lines) <= 1: parsed["help"] = help_text return parsed description = lines[1] arguments_part = None match = _ARGUMENT_META_REGEX.search(description) if match: arguments_part = description[match.start():] description = description[:match.start()].strip() else: description = description.strip() arguments = parsed["arguments"] if arguments_part: for argument_type, name, doc in _TYPE_PARAM_REGEX.findall(arguments_part): if name not in arguments: arguments[name] = {"names": [to_caterpillar_case(name)]} arguments[name]["required"] = False arguments[name]["help"] = doc if "type" not in arguments[name]: arguments[name] = _parse_type(locate(argument_type), arguments[name]) else: log.warning("Attempted to set type in both function annotation and docstring for argument %s" % name) for name, doc in _PARAM_REGEX.findall(arguments_part): if name not in arguments: arguments[name] = {"names": [to_caterpillar_case(name)]} arguments[name]["required"] = False arguments[name]["help"] = doc for name, aliases in _ALIAS_REGEX.findall(arguments_part): if name not in arguments: raise NameError("%s is not an argument. An argument has to exist to be aliased." % name) for alias in aliases.split(","): arguments[name]["names"].append(to_caterpillar_case(alias.strip())) log.debug( "%s: Finished parsing function metadata:\n\t%s\n\t%s\n\t%s" % (function.__name__, help_text, description, arguments) ) parsed["arguments"] = arguments parsed["help"] = help_text parsed["description"] = description return parsed def _parse_function(function: callable) -> dict: """Parses a functions parameters and help-text from its docstring and annotations. :param function: A function :return: The description and arguments (in order) for the function """ parsed = {} log.debug("Parsing arguments and docstring for %s..." % function.__name__) parsed["arguments"] = _parse_parameters(function) if not function.__doc__: return parsed parsed = _parse_docstring(function, parsed) return parsed
[docs]def register_plugin(plugin_sequence: tuple) -> None: """Creates a parser for a plugin :param plugin_sequence: The plugin sequence :raises KeyError: If the plugin already exists """ log.debug("Registering parser for plugin \"%s\"" % (plugin_sequence,)) _PARSERS.add_parser(plugin_sequence)
[docs]def register_command(plugin_sequence: tuple, command_name: str) -> None: """Registers a command towards a plugin. :param plugin_sequence: The name of the plugin that the function belongs to :param command_name: The name of the function """ log.debug("Registering parser for command \"%s\" in plugin \"%s\"" % (command_name, "/".join(plugin_sequence))) _PARSERS.add_command(plugin_sequence, command_name)
[docs]def register_command_metadata(plugin_sequence: tuple, command_name: str, function: callable) -> None: """Registers the function metadata to a parser Registering the metadata is done at a later point to not do unnecessary operations with _parse_function(). A single call to _parse_function() takes quite some time, so it's better to do it JIT. :param plugin_sequence: The name of the plugin :param command_name: The name of the command :param function: The function to parse metadata from """ log.debug("Registering metadata for \"%s\" on plugin \"%s\"" % (command_name, "/".join(plugin_sequence))) command = _PARSERS.get_command(plugin_sequence, command_name) metadata = _parse_function(function) if "help" in metadata and metadata["help"]: command.help = metadata["help"] if "description" in metadata and metadata["description"]: command.description = metadata["description"] for options in metadata["arguments"].values(): if "required" in options and not options["required"]: names = [] for name in options.pop("names"): prefix = "-" if len(name) is 1 else "--" names.append(prefix + name) else: names = options.pop("names") command.add_argument(*names, **options)
[docs]def get_plugin_arguments() -> tuple: """Only parse the plugin arguments. Plugin arguments are the <plugin> and <command> part of the CLI. If we do not do it like this help-messages for commands would be catched in here, and thus we wouldn't get a proper help message for commands. :returns: Plugin and command name """ log.debug("Parsing plugin arguments") base_args = [arg for arg in sys.argv[1:] if arg != "--help"] args = vars(_ROOT_PARSER.parse_known_args(base_args)[0]) plugin_name = tuple(args[_PLUGIN_ARGUMENT][:-1]) command_name = args[_PLUGIN_ARGUMENT][-1] return plugin_name, command_name
[docs]def get_command_arguments() -> dict: """Gets the command arguments. See get_plugin_arguments() for why we do it like this. :returns: Command arguments """ log.debug("Parsing command arguments") parsed = vars(_ROOT_PARSER.parse_args()) # Remove plugin arguments as they aren't necessary for the command arguments parsed.pop(_PLUGIN_ARGUMENT) return parsed