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",
)
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_plugin(self, plugin: tuple):
"""Gets a plugin parser"""
return self._dict[plugin]
[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 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
"""
def error_or_help(parser):
if any(keyword in sys.argv for keyword in ("--help", "-h")):
parser.print_help()
parser.exit()
parser.error("Not enough arguments")
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])
if _PLUGIN_ARGUMENT not in args:
error_or_help(_ROOT_PARSER)
plugin_name = tuple(args[_PLUGIN_ARGUMENT][:-1])
command_name = args[_PLUGIN_ARGUMENT][-1]
try:
parser = _PARSERS.get_command(plugin_name, command_name)
except KeyError:
parser = _PARSERS.get_plugin(command_name) # If there is just one argument it will end up here
if isinstance(parser, dict):
error_or_help(parser[_PLUGIN_PARSER])
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