diff options
| author | Eric Lin <anselor@gmail.com> | 2018-04-16 15:11:13 -0400 |
|---|---|---|
| committer | Eric Lin <anselor@gmail.com> | 2018-04-16 15:11:16 -0400 |
| commit | bb5e35803a491e85e91178e71fba9b362b08b3e7 (patch) | |
| tree | db7b531a8687d936189c873c4c6f4b691dda34cf | |
| parent | 21454b296867b0fce9f1afb70ca0476fe4b9ae2f (diff) | |
| download | cmd2-git-bb5e35803a491e85e91178e71fba9b362b08b3e7.tar.gz | |
Added more advanced/complex autocompleter examples.
Added more type hinting to AutoCompleter.
| -rwxr-xr-x | AutoCompleter.py | 1620 | ||||
| -rwxr-xr-x | examples/tab_autocompletion.py | 194 |
2 files changed, 987 insertions, 827 deletions
diff --git a/AutoCompleter.py b/AutoCompleter.py index 014ea8ab..e2af5ed9 100755 --- a/AutoCompleter.py +++ b/AutoCompleter.py @@ -1,807 +1,813 @@ -# coding=utf-8
-import argparse
-import re as _re
-import sys
-from argparse import OPTIONAL, ZERO_OR_MORE, ONE_OR_MORE, REMAINDER, PARSER, ArgumentError
-from typing import List, Dict, Tuple, Callable, Union
-
-from colorama import Fore
-
-from cmd2 import readline_lib
-
-
-class _RangeAction(object):
- def __init__(self, nargs):
- self.nargs_min = None
- self.nargs_max = None
-
- # pre-process special ranged nargs
- if isinstance(nargs, tuple):
- if len(nargs) != 2 or not isinstance(nargs[0], int) or not isinstance(nargs[1], int):
- raise ValueError('Ranged values for nargs must be a tuple of 2 integers')
- if nargs[0] >= nargs[1]:
- raise ValueError('Invalid nargs range. The first value must be less than the second')
- if nargs[0] < 0:
- raise ValueError('Negative numbers are invalid for nargs range.')
- narg_range = nargs
- self.nargs_min = nargs[0]
- self.nargs_max = nargs[1]
- if narg_range[0] == 0:
- if narg_range[1] > 1:
- self.nargs_adjusted = '*'
- else:
- # this shouldn't use a range tuple, but yet here we are
- self.nargs_adjusted = '?'
- else:
- self.nargs_adjusted = '+'
- else:
- self.nargs_adjusted = nargs
-
-
-class _StoreRangeAction(argparse._StoreAction, _RangeAction):
- def __init__(self,
- option_strings,
- dest,
- nargs=None,
- const=None,
- default=None,
- type=None,
- choices=None,
- required=False,
- help=None,
- metavar=None):
-
- _RangeAction.__init__(self, nargs)
-
- argparse._StoreAction.__init__(self,
- option_strings=option_strings,
- dest=dest,
- nargs=self.nargs_adjusted,
- const=const,
- default=default,
- type=type,
- choices=choices,
- required=required,
- help=help,
- metavar=metavar)
-
-
-class _AppendRangeAction(argparse._AppendAction, _RangeAction):
- def __init__(self,
- option_strings,
- dest,
- nargs=None,
- const=None,
- default=None,
- type=None,
- choices=None,
- required=False,
- help=None,
- metavar=None):
-
- _RangeAction.__init__(self, nargs)
-
- argparse._AppendAction.__init__(self,
- option_strings=option_strings,
- dest=dest,
- nargs=self.nargs_adjusted,
- const=const,
- default=default,
- type=type,
- choices=choices,
- required=required,
- help=help,
- metavar=metavar)
-
-
-def register_custom_actions(parser: argparse.ArgumentParser):
- """Register custom argument action types"""
- parser.register('action', None, _StoreRangeAction)
- parser.register('action', 'store', _StoreRangeAction)
- parser.register('action', 'append', _AppendRangeAction)
-
-
-class AutoCompleter(object):
- """Automatically command line tab completion based on argparse parameters"""
-
- class _ArgumentState(object):
- def __init__(self):
- self.min = None
- self.max = None
- self.count = 0
- self.needed = False
- self.variable = False
-
- def reset(self):
- """reset tracking values"""
- self.min = None
- self.max = None
- self.count = 0
- self.needed = False
- self.variable = False
-
- def __init__(self,
- parser: argparse.ArgumentParser,
- token_start_index: int = 1,
- arg_choices: Dict[str, Union[List, Tuple, Callable]] = None,
- subcmd_args_lookup: dict = None,
- tab_for_arg_help: bool = True):
- """
- AutoCompleter interprets the argparse.ArgumentParser internals to automatically
- generate the completion options for each argument.
-
- How to supply completion options for each argument:
- argparse Choices
- - pass a list of values to the choices parameter of an argparse argument.
- ex: parser.add_argument('-o', '--options', dest='options', choices=['An Option', 'SomeOtherOption'])
-
- arg_choices dictionary lookup
- arg_choices is a dict() mapping from argument name to one of 3 possible values:
- ex:
- parser = argparse.ArgumentParser()
- parser.add_argument('-o', '--options', dest='options')
- choices = {}
- mycompleter = AutoCompleter(parser, completer, 1, choices)
-
- - static list - provide a static list for each argument name
- ex:
- choices['options'] = ['An Option', 'SomeOtherOption']
-
- - choices function - provide a function that returns a list for each argument name
- ex:
- def generate_choices():
- return ['An Option', 'SomeOtherOption']
- choices['options'] = generate_choices
-
- - custom completer function - provide a completer function that will return the list
- of completion arguments
- ex 1:
- def my_completer(text: str, line: str, begidx: int, endidx:int):
- my_choices = [...]
- return my_choices
- choices['options'] = (my_completer)
- ex 2:
- def my_completer(text: str, line: str, begidx: int, endidx:int, extra_param: str, another: int):
- my_choices = [...]
- return my_choices
- completer_params = {'extra_param': 'my extra', 'another': 5}
- choices['options'] = (my_completer, completer_params)
-
- How to supply completion choice lists or functions for sub-commands:
- subcmd_args_lookup is used to supply a unique pair of arg_choices and subcmd_args_lookup
- for each sub-command in an argparser subparser group.
- This requires your subparser group to be named with the dest parameter
- ex:
- parser = ArgumentParser()
- subparsers = parser.add_subparsers(title='Actions', dest='action')
-
- subcmd_args_lookup maps a named subparser group to a subcommand group dictionary
- The subcommand group dictionary maps subcommand names to tuple(arg_choices, subcmd_args_lookup)
-
-
- :param parser: ArgumentParser instance
- :param token_start_index: index of the token to start parsing at
- :param arg_choices: dictionary mapping from argparse argument 'dest' name to list of choices
- :param subcmd_args_lookup: mapping a sub-command group name to a tuple to fill the child\
- AutoCompleter's arg_choices and subcmd_args_lookup parameters
- """
- if not subcmd_args_lookup:
- subcmd_args_lookup = {}
- forward_arg_choices = True
- else:
- forward_arg_choices = False
- self._parser = parser
- self._arg_choices = arg_choices.copy() if arg_choices is not None else {}
- self._token_start_index = token_start_index
- self._tab_for_arg_help = tab_for_arg_help
-
- self._flags = [] # all flags in this command
- self._flags_without_args = [] # all flags that don't take arguments
- self._flag_to_action = {} # maps flags to the argparse action object
- self._positional_actions = [] # argument names for positional arguments (by position index)
- self._positional_completers = {} # maps action name to sub-command autocompleter:
- # action_name -> dict(sub_command -> completer)
-
- # Start digging through the argparse structures.
- # _actions is the top level container of parameter definitions
- for action in self._parser._actions:
- # if there are choices defined, record them in the arguments dictionary
- if action.choices is not None:
- self._arg_choices[action.dest] = action.choices
-
- # if the parameter is flag based, it will have option_strings
- if action.option_strings:
- # record each option flag
- for option in action.option_strings:
- self._flags.append(option)
- self._flag_to_action[option] = action
- if action.nargs == 0:
- self._flags_without_args.append(option)
- else:
- self._positional_actions.append(action)
-
- if isinstance(action, argparse._SubParsersAction):
- sub_completers = {}
- sub_commands = []
- args_for_action = subcmd_args_lookup[action.dest] if action.dest in subcmd_args_lookup.keys() else {}
- for subcmd in action.choices:
- (subcmd_args, subcmd_lookup) = args_for_action[subcmd] if subcmd in args_for_action.keys() else (arg_choices, subcmd_args_lookup) if forward_arg_choices else ({}, {})
- subcmd_start = token_start_index + len(self._positional_actions)
- sub_completers[subcmd] = AutoCompleter(action.choices[subcmd], subcmd_start,
- arg_choices=subcmd_args, subcmd_args_lookup=subcmd_lookup)
- sub_commands.append(subcmd)
- self._positional_completers[action.dest] = sub_completers
- self._arg_choices[action.dest] = sub_commands
-
- def complete_command(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int):
- """Complete the command using the argparse metadata and provided argument dictionary"""
- # Count which positional argument index we're at now. Loop through all tokens on the command line so far
- # Skip any flags or flag parameter tokens
- next_pos_arg_index = 0
-
- pos_arg = AutoCompleter._ArgumentState()
- pos_action = None
-
- flag_arg = AutoCompleter._ArgumentState()
- flag_action = None
-
- matched_flags = []
- current_is_positional = False
- consumed_arg_values = {} # dict(arg_name -> [values, ...])
-
- def consume_flag_argument():
- """Consuming token as a flag argument"""
- # we're consuming flag arguments
- # if this is not empty and is not another potential flag, count towards flag arguments
- if len(token) > 0 and not token[0] in self._parser.prefix_chars and flag_action is not None:
- flag_arg.count += 1
-
- # does this complete a option item for the flag
- arg_choices = self._resolve_choices_for_arg(flag_action)
- # if the current token matches the current position's autocomplete argument list,
- # track that we've used it already. Unless this is the current token, then keep it.
- if not is_last_token and token in arg_choices:
- consumed_arg_values.setdefault(flag_action.dest, [])
- consumed_arg_values[flag_action.dest].append(token)
-
- def consume_positional_argument():
- """Consuming token as positional argument"""
- pos_arg.count += 1
-
- # does this complete a option item for the flag
- arg_choices = self._resolve_choices_for_arg(pos_action)
- # if the current token matches the current position's autocomplete argument list,
- # track that we've used it already. Unless this is the current token, then keep it.
- if not is_last_token and token in arg_choices:
- consumed_arg_values.setdefault(pos_action.dest, [])
- consumed_arg_values[pos_action.dest].append(token)
-
- is_last_token = False
- for idx, token in enumerate(tokens):
- is_last_token = idx >= len(tokens) - 1
- # Only start at the start token index
- if idx >= self._token_start_index:
- current_is_positional = False
- # Are we consuming flag arguments?
- if not flag_arg.needed:
- # we're not consuming flag arguments, is the current argument a potential flag?
- if len(token) > 0 and token[0] in self._parser.prefix_chars and\
- (is_last_token or (not is_last_token and token != '-')):
- # reset some tracking values
- flag_arg.reset()
- # don't reset positional tracking because flags can be interspersed anywhere between positionals
- flag_action = None
-
- # does the token fully match a known flag?
- if token in self._flag_to_action.keys():
- flag_action = self._flag_to_action[token]
- elif self._parser.allow_abbrev:
- candidates_flags = [flag for flag in self._flag_to_action.keys() if flag.startswith(token)]
- if len(candidates_flags) == 1:
- flag_action = self._flag_to_action[candidates_flags[0]]
-
- if flag_action is not None:
- # resolve argument counts
- self._process_action_nargs(flag_action, flag_arg)
- if not is_last_token and not isinstance(flag_action, argparse._AppendAction):
- matched_flags.extend(flag_action.option_strings)
-
- # current token isn't a potential flag
- # - does the last flag accept variable arguments?
- # - have we reached the max arg count for the flag?
- elif not flag_arg.variable or flag_arg.count >= flag_arg.max:
- # previous flag doesn't accept variable arguments, count this as a positional argument
-
- # reset flag tracking variables
- flag_arg.reset()
- flag_action = None
- current_is_positional = True
-
- if len(token) > 0 and pos_action is not None and pos_arg.count < pos_arg.max:
- # we have positional action match and we haven't reached the max arg count, consume
- # the positional argument and move on.
- consume_positional_argument()
- elif pos_action is None or pos_arg.count >= pos_arg.max:
- # if we don't have a current positional action or we've reached the max count for the action
- # close out the current positional argument state and set up for the next one
- pos_index = next_pos_arg_index
- next_pos_arg_index += 1
- pos_arg.reset()
- pos_action = None
-
- # are we at a sub-command? If so, forward to the matching completer
- if pos_index < len(self._positional_actions):
- action = self._positional_actions[pos_index]
- pos_name = action.dest
- if pos_name in self._positional_completers.keys():
- sub_completers = self._positional_completers[pos_name]
- if token in sub_completers.keys():
- return sub_completers[token].complete_command(tokens, text, line, begidx, endidx)
- pos_action = action
- self._process_action_nargs(pos_action, pos_arg)
- consume_positional_argument()
-
- elif not is_last_token and pos_arg.max is not None:
- pos_action = None
- pos_arg.reset()
-
- else:
- consume_flag_argument()
-
- else:
- consume_flag_argument()
-
- # don't reset this if we're on the last token - this allows completion to occur on the current token
- if not is_last_token and flag_arg.min is not None:
- flag_arg.needed = flag_arg.count < flag_arg.min
-
- # if we don't have a flag to populate with arguments and the last token starts with
- # a flag prefix then we'll complete the list of flag options
- completion_results = []
- if not flag_arg.needed and len(tokens[-1]) > 0 and tokens[-1][0] in self._parser.prefix_chars:
- return AutoCompleter.basic_complete(text, line, begidx, endidx,
- [flag for flag in self._flags if flag not in matched_flags])
- # we're not at a positional argument, see if we're in a flag argument
- elif not current_is_positional:
- # current_items = []
- if flag_action is not None:
- consumed = consumed_arg_values[flag_action.dest] if flag_action.dest in consumed_arg_values.keys() else []
- # current_items.extend(self._resolve_choices_for_arg(flag_action, consumed))
- completion_results = self._complete_for_arg(flag_action, text, line, begidx, endidx, consumed)
- if not completion_results:
- self._print_action_help(flag_action)
-
- # ok, we're not a flag, see if there's a positional argument to complete
- else:
- if pos_action is not None:
- pos_name = pos_action.dest
- consumed = consumed_arg_values[pos_name] if pos_name in consumed_arg_values.keys() else []
- completion_results = self._complete_for_arg(pos_action, text, line, begidx, endidx, consumed)
- if not completion_results:
- self._print_action_help(pos_action)
-
- return completion_results
-
- @staticmethod
- def _process_action_nargs(action, arg_state):
- if isinstance(action, _RangeAction):
- arg_state.min = action.nargs_min
- arg_state.max = action.nargs_max
- arg_state.variable = True
- if arg_state.min is None or arg_state.max is None:
- if action.nargs is None:
- arg_state.min = 1
- arg_state.max = 1
- elif action.nargs == '+':
- arg_state.min = 1
- arg_state.max = float('inf')
- arg_state.variable = True
- elif action.nargs == '*':
- arg_state.min = 0
- arg_state.max = float('inf')
- arg_state.variable = True
- elif action.nargs == '?':
- arg_state.min = 0
- arg_state.max = 1
- arg_state.variable = True
- else:
- arg_state.min = action.nargs
- arg_state.max = action.nargs
-
- def _complete_for_arg(self, action, text: str, line: str, begidx: int, endidx: int, used_values=list()):
- if action.dest in self._arg_choices.keys():
- arg_choices = self._arg_choices[action.dest]
-
- if isinstance(arg_choices, tuple) and len(arg_choices) > 0 and callable(arg_choices[0]):
- completer = arg_choices[0]
- list_args = None
- kw_args = None
- for index in range(1, len(arg_choices)):
- if isinstance(arg_choices[index], list):
- list_args = arg_choices[index]
- elif isinstance(arg_choices[index], dict):
- kw_args = arg_choices[index]
- if list_args is not None and kw_args is not None:
- return completer(text, line, begidx, endidx, *list_args, **kw_args)
- elif list_args is not None:
- return completer(text, line, begidx, endidx, *list_args)
- elif kw_args is not None:
- return completer(text, line, begidx, endidx, **kw_args)
- else:
- return completer(text, line, begidx, endidx)
- else:
- return AutoCompleter.basic_complete(text, line, begidx, endidx,
- self._resolve_choices_for_arg(action, used_values))
-
- return []
-
- def _resolve_choices_for_arg(self, action, used_values=list()):
- if action.dest in self._arg_choices.keys():
- args = self._arg_choices[action.dest]
- if callable(args):
- args = args()
-
- try:
- iter(args)
- except TypeError:
- pass
- else:
- # filter out arguments we already used
- args = [arg for arg in args if arg not in used_values]
-
- if len(args) > 0:
- return args
-
- return []
-
- def _print_action_help(self, action):
- if not self._tab_for_arg_help:
- return
- if action.option_strings:
- flags = ', '.join(action.option_strings)
- param = ''
- if action.nargs is None or action.nargs != 0:
- param += ' ' + str(action.dest).upper()
-
- prefix = '{}{}'.format(flags, param)
- else:
- prefix = '{}'.format(str(action.dest).upper())
-
- prefix = ' {0: <{width}} '.format(prefix, width=20)
- pref_len = len(prefix)
- help_lines = action.help.splitlines()
- if len(help_lines) == 1:
- print('\nHint:\n{}{}\n'.format(prefix, help_lines[0]))
- else:
- out_str = '\n{}'.format(prefix)
- out_str += '\n{0: <{width}}'.format('', width=pref_len).join(help_lines)
- print('\nHint:' + out_str + '\n')
-
- readline_lib.rl_forced_update_display()
-
- # noinspection PyUnusedLocal
- @staticmethod
- def basic_complete(text, line, begidx, endidx, match_against):
- """
- Performs tab completion against a list
-
- :param text: str - the string prefix we are attempting to match (all returned matches must begin with it)
- :param line: str - the current input line with leading whitespace removed
- :param begidx: int - the beginning index of the prefix text
- :param endidx: int - the ending index of the prefix text
- :param match_against: Collection - the list being matched against
- :return: List[str] - a list of possible tab completions
- """
- return [cur_match for cur_match in match_against if cur_match.startswith(text)]
-
-
-# Copied from argparse
-from argparse import _
-
-
-class ACHelpFormatter(argparse.HelpFormatter):
- """Custom help formatter to configure ordering of help text"""
-
- def _format_usage(self, usage, actions, groups, prefix):
- if prefix is None:
- prefix = _('Usage: ')
-
- # if usage is specified, use that
- if usage is not None:
- usage = usage % dict(prog=self._prog)
-
- # if no optionals or positionals are available, usage is just prog
- elif usage is None and not actions:
- usage = '%(prog)s' % dict(prog=self._prog)
-
- # if optionals and positionals are available, calculate usage
- elif usage is None:
- prog = '%(prog)s' % dict(prog=self._prog)
-
- # split optionals from positionals
- optionals = []
- required_options = []
- positionals = []
- for action in actions:
- if action.option_strings:
- if action.required:
- required_options.append(action)
- else:
- optionals.append(action)
- else:
- positionals.append(action)
-
- # build full usage string
- format = self._format_actions_usage
- action_usage = format(positionals + required_options + optionals, groups)
- usage = ' '.join([s for s in [prog, action_usage] if s])
-
- # wrap the usage parts if it's too long
- text_width = self._width - self._current_indent
- if len(prefix) + len(usage) > text_width:
-
- # break usage into wrappable parts
- part_regexp = r'\(.*?\)+|\[.*?\]+|\S+'
- opt_usage = format(optionals, groups)
- pos_usage = format(positionals, groups)
- req_usage = format(required_options, groups)
- opt_parts = _re.findall(part_regexp, opt_usage)
- pos_parts = _re.findall(part_regexp, pos_usage)
- req_parts = _re.findall(part_regexp, req_usage)
- assert ' '.join(opt_parts) == opt_usage
- assert ' '.join(pos_parts) == pos_usage
- assert ' '.join(req_parts) == req_usage
-
- # helper for wrapping lines
- def get_lines(parts, indent, prefix=None):
- lines = []
- line = []
- if prefix is not None:
- line_len = len(prefix) - 1
- else:
- line_len = len(indent) - 1
- for part in parts:
- if line_len + 1 + len(part) > text_width and line:
- lines.append(indent + ' '.join(line))
- line = []
- line_len = len(indent) - 1
- line.append(part)
- line_len += len(part) + 1
- if line:
- lines.append(indent + ' '.join(line))
- if prefix is not None:
- lines[0] = lines[0][len(indent):]
- return lines
-
- # if prog is short, follow it with optionals or positionals
- if len(prefix) + len(prog) <= 0.75 * text_width:
- indent = ' ' * (len(prefix) + len(prog) + 1)
- if opt_parts:
- lines = get_lines([prog] + pos_parts, indent, prefix)
- lines.extend(get_lines(req_parts, indent))
- lines.extend(get_lines(opt_parts, indent))
- elif pos_parts:
- lines = get_lines([prog] + pos_parts, indent, prefix)
- lines.extend(get_lines(req_parts, indent))
- else:
- lines = [prog]
-
- # if prog is long, put it on its own line
- else:
- indent = ' ' * len(prefix)
- parts = pos_parts + req_parts + opt_parts
- lines = get_lines(parts, indent)
- if len(lines) > 1:
- lines = []
- lines.extend(get_lines(pos_parts, indent))
- lines.extend(get_lines(req_parts, indent))
- lines.extend(get_lines(opt_parts, indent))
- lines = [prog] + lines
-
- # join lines into usage
- usage = '\n'.join(lines)
-
- # prefix with 'usage:'
- return '%s%s\n\n' % (prefix, usage)
-
- def _format_action_invocation(self, action):
- if not action.option_strings:
- default = self._get_default_metavar_for_positional(action)
- metavar, = self._metavar_formatter(action, default)(1)
- return metavar
-
- else:
- parts = []
-
- # if the Optional doesn't take a value, format is:
- # -s, --long
- if action.nargs == 0:
- parts.extend(action.option_strings)
- return ', '.join(parts)
-
- # if the Optional takes a value, format is:
- # -s ARGS, --long ARGS
- else:
- default = self._get_default_metavar_for_optional(action)
- args_string = self._format_args(action, default)
-
- # for option_string in action.option_strings:
- # parts.append('%s %s' % (option_string, args_string))
-
- return ', '.join(action.option_strings) + ' ' + args_string
-
- def _format_args(self, action, default_metavar):
- get_metavar = self._metavar_formatter(action, default_metavar)
- if isinstance(action, _RangeAction) and \
- action.nargs_min is not None and action.nargs_max is not None:
- result = '{}{{{}..{}}}'.format('%s' % get_metavar(1), action.nargs_min, action.nargs_max)
-
- elif action.nargs is None:
- result = '%s' % get_metavar(1)
- elif action.nargs == OPTIONAL:
- result = '[%s]' % get_metavar(1)
- elif action.nargs == ZERO_OR_MORE:
- result = '[%s [...]]' % get_metavar(1)
- elif action.nargs == ONE_OR_MORE:
- result = '%s [...]' % get_metavar(1)
- elif action.nargs == REMAINDER:
- result = '...'
- elif action.nargs == PARSER:
- result = '%s ...' % get_metavar(1)
- else:
- formats = ['%s' for _ in range(action.nargs)]
- result = ' '.join(formats) % get_metavar(action.nargs)
- return result
-
- def _metavar_formatter(self, action, default_metavar):
- if action.metavar is not None:
- result = action.metavar
- elif action.choices is not None:
- choice_strs = [str(choice) for choice in action.choices]
- result = '{%s}' % ', '.join(choice_strs)
- else:
- result = default_metavar
-
- def format(tuple_size):
- if isinstance(result, tuple):
- return result
- else:
- return (result, ) * tuple_size
- return format
-
- def _split_lines(self, text, width):
- return text.splitlines()
-
-
-class ACArgumentParser(argparse.ArgumentParser):
- """Custom argparse class to override error method to change default help text."""
-
- def __init__(self,
- prog=None,
- usage=None,
- description=None,
- epilog=None,
- parents=[],
- formatter_class=ACHelpFormatter,
- prefix_chars='-',
- fromfile_prefix_chars=None,
- argument_default=None,
- conflict_handler='error',
- add_help=True,
- allow_abbrev=True):
-
- super().__init__(prog=prog,
- usage=usage,
- description=description,
- epilog=epilog,
- parents=parents,
- formatter_class=formatter_class,
- prefix_chars=prefix_chars,
- fromfile_prefix_chars=fromfile_prefix_chars,
- argument_default=argument_default,
- conflict_handler=conflict_handler,
- add_help=add_help,
- allow_abbrev=allow_abbrev)
- register_custom_actions(self)
-
- self._custom_error_message = ''
-
- def set_custom_message(self, custom_message=''):
- """
- Allows an error message override to the error() function, useful when forcing a
- re-parse of arguments with newly required parameters
- """
- self._custom_error_message = custom_message
-
- def error(self, message):
- """Custom error override."""
- if len(self._custom_error_message) > 0:
- message = self._custom_error_message
- self._custom_error_message = ''
-
- lines = message.split('\n')
- linum = 0
- formatted_message = ''
- for line in lines:
- if linum == 0:
- formatted_message = 'Error: ' + line
- else:
- formatted_message += '\n ' + line
- linum += 1
-
- sys.stderr.write(Fore.LIGHTRED_EX + '{}\n\n'.format(formatted_message) + Fore.RESET)
- self.print_help()
- sys.exit(1)
-
- def format_help(self):
- """Copy of format_help() from argparse.ArgumentParser with tweaks to separately display required parameters"""
- formatter = self._get_formatter()
-
- # usage
- formatter.add_usage(self.usage, self._actions,
- self._mutually_exclusive_groups)
-
- # description
- formatter.add_text(self.description)
-
- # positionals, optionals and user-defined groups
- for action_group in self._action_groups:
- if action_group.title == 'optional arguments':
- # check if the arguments are required, group accordingly
- req_args = []
- opt_args = []
- for action in action_group._group_actions:
- if action.required:
- req_args.append(action)
- else:
- opt_args.append(action)
-
- # separately display required arguments
- formatter.start_section('required arguments')
- formatter.add_text(action_group.description)
- formatter.add_arguments(req_args)
- formatter.end_section()
-
- # now display truly optional arguments
- formatter.start_section(action_group.title)
- formatter.add_text(action_group.description)
- formatter.add_arguments(opt_args)
- formatter.end_section()
- else:
- formatter.start_section(action_group.title)
- formatter.add_text(action_group.description)
- formatter.add_arguments(action_group._group_actions)
- formatter.end_section()
-
- # epilog
- formatter.add_text(self.epilog)
-
- # determine help from format above
- return formatter.format_help()
-
- def _get_nargs_pattern(self, action):
- # Override _get_nargs_pattern behavior to use the nargs ranges provided by AutoCompleter
- if isinstance(action, _RangeAction) and \
- action.nargs_min is not None and action.nargs_max is not None:
- nargs_pattern = '(-*A{{{},{}}}-*)'.format(action.nargs_min, action.nargs_max)
-
- # if this is an optional action, -- is not allowed
- if action.option_strings:
- nargs_pattern = nargs_pattern.replace('-*', '')
- nargs_pattern = nargs_pattern.replace('-', '')
- return nargs_pattern
- return super(ACArgumentParser, self)._get_nargs_pattern(action)
-
- def _match_argument(self, action, arg_strings_pattern):
- # match the pattern for this action to the arg strings
- nargs_pattern = self._get_nargs_pattern(action)
- match = _re.match(nargs_pattern, arg_strings_pattern)
-
- # raise an exception if we weren't able to find a match
- if match is None:
- if isinstance(action, _RangeAction) and \
- action.nargs_min is not None and action.nargs_max is not None:
- raise ArgumentError(action,
- 'Expected between {} and {} arguments'.format(action.nargs_min, action.nargs_max))
-
- return super(ACArgumentParser, self)._match_argument(action, arg_strings_pattern)
+# coding=utf-8 +import argparse +import re as _re +import sys +from argparse import OPTIONAL, ZERO_OR_MORE, ONE_OR_MORE, REMAINDER, PARSER, ArgumentError, _ +from typing import List, Dict, Tuple, Callable, Union + +from colorama import Fore + + +class _RangeAction(object): + def __init__(self, nargs: Union[int, str, Tuple[int, int], None]): + self.nargs_min = None + self.nargs_max = None + + # pre-process special ranged nargs + if isinstance(nargs, tuple): + if len(nargs) != 2 or not isinstance(nargs[0], int) or not isinstance(nargs[1], int): + raise ValueError('Ranged values for nargs must be a tuple of 2 integers') + if nargs[0] >= nargs[1]: + raise ValueError('Invalid nargs range. The first value must be less than the second') + if nargs[0] < 0: + raise ValueError('Negative numbers are invalid for nargs range.') + narg_range = nargs + self.nargs_min = nargs[0] + self.nargs_max = nargs[1] + if narg_range[0] == 0: + if narg_range[1] > 1: + self.nargs_adjusted = '*' + else: + # this shouldn't use a range tuple, but yet here we are + self.nargs_adjusted = '?' + else: + self.nargs_adjusted = '+' + else: + self.nargs_adjusted = nargs + + +class _StoreRangeAction(argparse._StoreAction, _RangeAction): + def __init__(self, + option_strings, + dest, + nargs=None, + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None): + + _RangeAction.__init__(self, nargs) + + argparse._StoreAction.__init__(self, + option_strings=option_strings, + dest=dest, + nargs=self.nargs_adjusted, + const=const, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar) + + +class _AppendRangeAction(argparse._AppendAction, _RangeAction): + def __init__(self, + option_strings, + dest, + nargs=None, + const=None, + default=None, + type=None, + choices=None, + required=False, + help=None, + metavar=None): + + _RangeAction.__init__(self, nargs) + + argparse._AppendAction.__init__(self, + option_strings=option_strings, + dest=dest, + nargs=self.nargs_adjusted, + const=const, + default=default, + type=type, + choices=choices, + required=required, + help=help, + metavar=metavar) + + +def register_custom_actions(parser: argparse.ArgumentParser): + """Register custom argument action types""" + parser.register('action', None, _StoreRangeAction) + parser.register('action', 'store', _StoreRangeAction) + parser.register('action', 'append', _AppendRangeAction) + + +class AutoCompleter(object): + """Automatically command line tab completion based on argparse parameters""" + + class _ArgumentState(object): + def __init__(self): + self.min = None + self.max = None + self.count = 0 + self.needed = False + self.variable = False + + def reset(self): + """reset tracking values""" + self.min = None + self.max = None + self.count = 0 + self.needed = False + self.variable = False + + def __init__(self, + parser: argparse.ArgumentParser, + token_start_index: int = 1, + arg_choices: Dict[str, Union[List, Tuple, Callable]] = None, + subcmd_args_lookup: dict = None, + tab_for_arg_help: bool = True): + """ + AutoCompleter interprets the argparse.ArgumentParser internals to automatically + generate the completion options for each argument. + + How to supply completion options for each argument: + argparse Choices + - pass a list of values to the choices parameter of an argparse argument. + ex: parser.add_argument('-o', '--options', dest='options', choices=['An Option', 'SomeOtherOption']) + + arg_choices dictionary lookup + arg_choices is a dict() mapping from argument name to one of 3 possible values: + ex: + parser = argparse.ArgumentParser() + parser.add_argument('-o', '--options', dest='options') + choices = {} + mycompleter = AutoCompleter(parser, completer, 1, choices) + + - static list - provide a static list for each argument name + ex: + choices['options'] = ['An Option', 'SomeOtherOption'] + + - choices function - provide a function that returns a list for each argument name + ex: + def generate_choices(): + return ['An Option', 'SomeOtherOption'] + choices['options'] = generate_choices + + - custom completer function - provide a completer function that will return the list + of completion arguments + ex 1: + def my_completer(text: str, line: str, begidx: int, endidx:int): + my_choices = [...] + return my_choices + choices['options'] = (my_completer) + ex 2: + def my_completer(text: str, line: str, begidx: int, endidx:int, extra_param: str, another: int): + my_choices = [...] + return my_choices + completer_params = {'extra_param': 'my extra', 'another': 5} + choices['options'] = (my_completer, completer_params) + + How to supply completion choice lists or functions for sub-commands: + subcmd_args_lookup is used to supply a unique pair of arg_choices and subcmd_args_lookup + for each sub-command in an argparser subparser group. + This requires your subparser group to be named with the dest parameter + ex: + parser = ArgumentParser() + subparsers = parser.add_subparsers(title='Actions', dest='action') + + subcmd_args_lookup maps a named subparser group to a subcommand group dictionary + The subcommand group dictionary maps subcommand names to tuple(arg_choices, subcmd_args_lookup) + + + :param parser: ArgumentParser instance + :param token_start_index: index of the token to start parsing at + :param arg_choices: dictionary mapping from argparse argument 'dest' name to list of choices + :param subcmd_args_lookup: mapping a sub-command group name to a tuple to fill the child\ + AutoCompleter's arg_choices and subcmd_args_lookup parameters + """ + if not subcmd_args_lookup: + subcmd_args_lookup = {} + forward_arg_choices = True + else: + forward_arg_choices = False + self._parser = parser + self._arg_choices = arg_choices.copy() if arg_choices is not None else {} + self._token_start_index = token_start_index + self._tab_for_arg_help = tab_for_arg_help + + self._flags = [] # all flags in this command + self._flags_without_args = [] # all flags that don't take arguments + self._flag_to_action = {} # maps flags to the argparse action object + self._positional_actions = [] # argument names for positional arguments (by position index) + self._positional_completers = {} # maps action name to sub-command autocompleter: + # action_name -> dict(sub_command -> completer) + + # Start digging through the argparse structures. + # _actions is the top level container of parameter definitions + for action in self._parser._actions: + # if there are choices defined, record them in the arguments dictionary + if action.choices is not None: + self._arg_choices[action.dest] = action.choices + + # if the parameter is flag based, it will have option_strings + if action.option_strings: + # record each option flag + for option in action.option_strings: + self._flags.append(option) + self._flag_to_action[option] = action + if action.nargs == 0: + self._flags_without_args.append(option) + else: + self._positional_actions.append(action) + + if isinstance(action, argparse._SubParsersAction): + sub_completers = {} + sub_commands = [] + args_for_action = subcmd_args_lookup[action.dest]\ + if action.dest in subcmd_args_lookup.keys() else {} + for subcmd in action.choices: + (subcmd_args, subcmd_lookup) = args_for_action[subcmd] if \ + subcmd in args_for_action.keys() else \ + (arg_choices, subcmd_args_lookup) if forward_arg_choices else ({}, {}) + subcmd_start = token_start_index + len(self._positional_actions) + sub_completers[subcmd] = AutoCompleter(action.choices[subcmd], subcmd_start, + arg_choices=subcmd_args, + subcmd_args_lookup=subcmd_lookup) + sub_commands.append(subcmd) + self._positional_completers[action.dest] = sub_completers + self._arg_choices[action.dest] = sub_commands + + def complete_command(self, tokens: List[str], text: str, line: str, begidx: int, endidx: int) -> List[str]: + """Complete the command using the argparse metadata and provided argument dictionary""" + # Count which positional argument index we're at now. Loop through all tokens on the command line so far + # Skip any flags or flag parameter tokens + next_pos_arg_index = 0 + + pos_arg = AutoCompleter._ArgumentState() + pos_action = None + + flag_arg = AutoCompleter._ArgumentState() + flag_action = None + + matched_flags = [] + current_is_positional = False + consumed_arg_values = {} # dict(arg_name -> [values, ...]) + + def consume_flag_argument() -> None: + """Consuming token as a flag argument""" + # we're consuming flag arguments + # if this is not empty and is not another potential flag, count towards flag arguments + if len(token) > 0 and not token[0] in self._parser.prefix_chars and flag_action is not None: + flag_arg.count += 1 + + # does this complete a option item for the flag + arg_choices = self._resolve_choices_for_arg(flag_action) + # if the current token matches the current position's autocomplete argument list, + # track that we've used it already. Unless this is the current token, then keep it. + if not is_last_token and token in arg_choices: + consumed_arg_values.setdefault(flag_action.dest, []) + consumed_arg_values[flag_action.dest].append(token) + + def consume_positional_argument() -> None: + """Consuming token as positional argument""" + pos_arg.count += 1 + + # does this complete a option item for the flag + arg_choices = self._resolve_choices_for_arg(pos_action) + # if the current token matches the current position's autocomplete argument list, + # track that we've used it already. Unless this is the current token, then keep it. + if not is_last_token and token in arg_choices: + consumed_arg_values.setdefault(pos_action.dest, []) + consumed_arg_values[pos_action.dest].append(token) + + is_last_token = False + for idx, token in enumerate(tokens): + is_last_token = idx >= len(tokens) - 1 + # Only start at the start token index + if idx >= self._token_start_index: + current_is_positional = False + # Are we consuming flag arguments? + if not flag_arg.needed: + # we're not consuming flag arguments, is the current argument a potential flag? + if len(token) > 0 and token[0] in self._parser.prefix_chars and\ + (is_last_token or (not is_last_token and token != '-')): + # reset some tracking values + flag_arg.reset() + # don't reset positional tracking because flags can be interspersed anywhere between positionals + flag_action = None + + # does the token fully match a known flag? + if token in self._flag_to_action.keys(): + flag_action = self._flag_to_action[token] + elif self._parser.allow_abbrev: + candidates_flags = [flag for flag in self._flag_to_action.keys() if flag.startswith(token)] + if len(candidates_flags) == 1: + flag_action = self._flag_to_action[candidates_flags[0]] + + if flag_action is not None: + # resolve argument counts + self._process_action_nargs(flag_action, flag_arg) + if not is_last_token and not isinstance(flag_action, argparse._AppendAction): + matched_flags.extend(flag_action.option_strings) + + # current token isn't a potential flag + # - does the last flag accept variable arguments? + # - have we reached the max arg count for the flag? + elif not flag_arg.variable or flag_arg.count >= flag_arg.max: + # previous flag doesn't accept variable arguments, count this as a positional argument + + # reset flag tracking variables + flag_arg.reset() + flag_action = None + current_is_positional = True + + if len(token) > 0 and pos_action is not None and pos_arg.count < pos_arg.max: + # we have positional action match and we haven't reached the max arg count, consume + # the positional argument and move on. + consume_positional_argument() + elif pos_action is None or pos_arg.count >= pos_arg.max: + # if we don't have a current positional action or we've reached the max count for the action + # close out the current positional argument state and set up for the next one + pos_index = next_pos_arg_index + next_pos_arg_index += 1 + pos_arg.reset() + pos_action = None + + # are we at a sub-command? If so, forward to the matching completer + if pos_index < len(self._positional_actions): + action = self._positional_actions[pos_index] + pos_name = action.dest + if pos_name in self._positional_completers.keys(): + sub_completers = self._positional_completers[pos_name] + if token in sub_completers.keys(): + return sub_completers[token].complete_command(tokens, text, line, + begidx, endidx) + pos_action = action + self._process_action_nargs(pos_action, pos_arg) + consume_positional_argument() + + elif not is_last_token and pos_arg.max is not None: + pos_action = None + pos_arg.reset() + + else: + consume_flag_argument() + + else: + consume_flag_argument() + + # don't reset this if we're on the last token - this allows completion to occur on the current token + if not is_last_token and flag_arg.min is not None: + flag_arg.needed = flag_arg.count < flag_arg.min + + # if we don't have a flag to populate with arguments and the last token starts with + # a flag prefix then we'll complete the list of flag options + completion_results = [] + if not flag_arg.needed and len(tokens[-1]) > 0 and tokens[-1][0] in self._parser.prefix_chars: + return AutoCompleter.basic_complete(text, line, begidx, endidx, + [flag for flag in self._flags if flag not in matched_flags]) + # we're not at a positional argument, see if we're in a flag argument + elif not current_is_positional: + # current_items = [] + if flag_action is not None: + consumed = consumed_arg_values[flag_action.dest]\ + if flag_action.dest in consumed_arg_values.keys() else [] + # current_items.extend(self._resolve_choices_for_arg(flag_action, consumed)) + completion_results = self._complete_for_arg(flag_action, text, line, begidx, endidx, consumed) + if not completion_results: + self._print_action_help(flag_action) + + # ok, we're not a flag, see if there's a positional argument to complete + else: + if pos_action is not None: + pos_name = pos_action.dest + consumed = consumed_arg_values[pos_name] if pos_name in consumed_arg_values.keys() else [] + completion_results = self._complete_for_arg(pos_action, text, line, begidx, endidx, consumed) + if not completion_results: + self._print_action_help(pos_action) + + return completion_results + + @staticmethod + def _process_action_nargs(action: argparse.Action, arg_state: _ArgumentState) -> None: + if isinstance(action, _RangeAction): + arg_state.min = action.nargs_min + arg_state.max = action.nargs_max + arg_state.variable = True + if arg_state.min is None or arg_state.max is None: + if action.nargs is None: + arg_state.min = 1 + arg_state.max = 1 + elif action.nargs == '+': + arg_state.min = 1 + arg_state.max = float('inf') + arg_state.variable = True + elif action.nargs == '*': + arg_state.min = 0 + arg_state.max = float('inf') + arg_state.variable = True + elif action.nargs == '?': + arg_state.min = 0 + arg_state.max = 1 + arg_state.variable = True + else: + arg_state.min = action.nargs + arg_state.max = action.nargs + + def _complete_for_arg(self, action: argparse.Action, + text: str, + line: str, + begidx: int, + endidx: int, + used_values=list()) -> List[str]: + if action.dest in self._arg_choices.keys(): + arg_choices = self._arg_choices[action.dest] + + if isinstance(arg_choices, tuple) and len(arg_choices) > 0 and callable(arg_choices[0]): + completer = arg_choices[0] + list_args = None + kw_args = None + for index in range(1, len(arg_choices)): + if isinstance(arg_choices[index], list): + list_args = arg_choices[index] + elif isinstance(arg_choices[index], dict): + kw_args = arg_choices[index] + if list_args is not None and kw_args is not None: + return completer(text, line, begidx, endidx, *list_args, **kw_args) + elif list_args is not None: + return completer(text, line, begidx, endidx, *list_args) + elif kw_args is not None: + return completer(text, line, begidx, endidx, **kw_args) + else: + return completer(text, line, begidx, endidx) + else: + return AutoCompleter.basic_complete(text, line, begidx, endidx, + self._resolve_choices_for_arg(action, used_values)) + + return [] + + def _resolve_choices_for_arg(self, action: argparse.Action, used_values=list()) -> List[str]: + if action.dest in self._arg_choices.keys(): + args = self._arg_choices[action.dest] + if callable(args): + args = args() + + try: + iter(args) + except TypeError: + pass + else: + # filter out arguments we already used + args = [arg for arg in args if arg not in used_values] + + if len(args) > 0: + return args + + return [] + + def _print_action_help(self, action: argparse.Action) -> None: + if not self._tab_for_arg_help: + return + if action.option_strings: + flags = ', '.join(action.option_strings) + param = '' + if action.nargs is None or action.nargs != 0: + param += ' ' + str(action.dest).upper() + + prefix = '{}{}'.format(flags, param) + else: + prefix = '{}'.format(str(action.dest).upper()) + + prefix = ' {0: <{width}} '.format(prefix, width=20) + pref_len = len(prefix) + help_lines = action.help.splitlines() + if len(help_lines) == 1: + print('\nHint:\n{}{}\n'.format(prefix, help_lines[0])) + else: + out_str = '\n{}'.format(prefix) + out_str += '\n{0: <{width}}'.format('', width=pref_len).join(help_lines) + print('\nHint:' + out_str + '\n') + + from cmd2 import readline_lib + readline_lib.rl_forced_update_display() + + # noinspection PyUnusedLocal + @staticmethod + def basic_complete(text: str, line: str, begidx: int, endidx: int, match_against: List[str]) -> List[str]: + """ + Performs tab completion against a list + + :param text: str - the string prefix we are attempting to match (all returned matches must begin with it) + :param line: str - the current input line with leading whitespace removed + :param begidx: int - the beginning index of the prefix text + :param endidx: int - the ending index of the prefix text + :param match_against: Collection - the list being matched against + :return: List[str] - a list of possible tab completions + """ + return [cur_match for cur_match in match_against if cur_match.startswith(text)] + + +class ACHelpFormatter(argparse.HelpFormatter): + """Custom help formatter to configure ordering of help text""" + + def _format_usage(self, usage, actions, groups, prefix): + if prefix is None: + prefix = _('Usage: ') + + # if usage is specified, use that + if usage is not None: + usage %= dict(prog=self._prog) + + # if no optionals or positionals are available, usage is just prog + elif usage is None and not actions: + usage = '%(prog)s' % dict(prog=self._prog) + + # if optionals and positionals are available, calculate usage + elif usage is None: + prog = '%(prog)s' % dict(prog=self._prog) + + # split optionals from positionals + optionals = [] + required_options = [] + positionals = [] + for action in actions: + if action.option_strings: + if action.required: + required_options.append(action) + else: + optionals.append(action) + else: + positionals.append(action) + + # build full usage string + format = self._format_actions_usage + action_usage = format(positionals + required_options + optionals, groups) + usage = ' '.join([s for s in [prog, action_usage] if s]) + + # wrap the usage parts if it's too long + text_width = self._width - self._current_indent + if len(prefix) + len(usage) > text_width: + + # break usage into wrappable parts + part_regexp = r'\(.*?\)+|\[.*?\]+|\S+' + opt_usage = format(optionals, groups) + pos_usage = format(positionals, groups) + req_usage = format(required_options, groups) + opt_parts = _re.findall(part_regexp, opt_usage) + pos_parts = _re.findall(part_regexp, pos_usage) + req_parts = _re.findall(part_regexp, req_usage) + assert ' '.join(opt_parts) == opt_usage + assert ' '.join(pos_parts) == pos_usage + assert ' '.join(req_parts) == req_usage + + # helper for wrapping lines + def get_lines(parts, indent, prefix=None): + lines = [] + line = [] + if prefix is not None: + line_len = len(prefix) - 1 + else: + line_len = len(indent) - 1 + for part in parts: + if line_len + 1 + len(part) > text_width and line: + lines.append(indent + ' '.join(line)) + line = [] + line_len = len(indent) - 1 + line.append(part) + line_len += len(part) + 1 + if line: + lines.append(indent + ' '.join(line)) + if prefix is not None: + lines[0] = lines[0][len(indent):] + return lines + + # if prog is short, follow it with optionals or positionals + if len(prefix) + len(prog) <= 0.75 * text_width: + indent = ' ' * (len(prefix) + len(prog) + 1) + if opt_parts: + lines = get_lines([prog] + pos_parts, indent, prefix) + lines.extend(get_lines(req_parts, indent)) + lines.extend(get_lines(opt_parts, indent)) + elif pos_parts: + lines = get_lines([prog] + pos_parts, indent, prefix) + lines.extend(get_lines(req_parts, indent)) + else: + lines = [prog] + + # if prog is long, put it on its own line + else: + indent = ' ' * len(prefix) + parts = pos_parts + req_parts + opt_parts + lines = get_lines(parts, indent) + if len(lines) > 1: + lines = [] + lines.extend(get_lines(pos_parts, indent)) + lines.extend(get_lines(req_parts, indent)) + lines.extend(get_lines(opt_parts, indent)) + lines = [prog] + lines + + # join lines into usage + usage = '\n'.join(lines) + + # prefix with 'usage:' + return '%s%s\n\n' % (prefix, usage) + + def _format_action_invocation(self, action): + if not action.option_strings: + default = self._get_default_metavar_for_positional(action) + metavar, = self._metavar_formatter(action, default)(1) + return metavar + + else: + parts = [] + + # if the Optional doesn't take a value, format is: + # -s, --long + if action.nargs == 0: + parts.extend(action.option_strings) + return ', '.join(parts) + + # if the Optional takes a value, format is: + # -s ARGS, --long ARGS + else: + default = self._get_default_metavar_for_optional(action) + args_string = self._format_args(action, default) + + # for option_string in action.option_strings: + # parts.append('%s %s' % (option_string, args_string)) + + return ', '.join(action.option_strings) + ' ' + args_string + + def _format_args(self, action, default_metavar): + get_metavar = self._metavar_formatter(action, default_metavar) + if isinstance(action, _RangeAction) and \ + action.nargs_min is not None and action.nargs_max is not None: + result = '{}{{{}..{}}}'.format('%s' % get_metavar(1), action.nargs_min, action.nargs_max) + + elif action.nargs is None: + result = '%s' % get_metavar(1) + elif action.nargs == OPTIONAL: + result = '[%s]' % get_metavar(1) + elif action.nargs == ZERO_OR_MORE: + result = '[%s [...]]' % get_metavar(1) + elif action.nargs == ONE_OR_MORE: + result = '%s [...]' % get_metavar(1) + elif action.nargs == REMAINDER: + result = '...' + elif action.nargs == PARSER: + result = '%s ...' % get_metavar(1) + else: + formats = ['%s' for _ in range(action.nargs)] + result = ' '.join(formats) % get_metavar(action.nargs) + return result + + def _metavar_formatter(self, action, default_metavar): + if action.metavar is not None: + result = action.metavar + elif action.choices is not None: + choice_strs = [str(choice) for choice in action.choices] + result = '{%s}' % ', '.join(choice_strs) + else: + result = default_metavar + + def format(tuple_size): + if isinstance(result, tuple): + return result + else: + return (result, ) * tuple_size + return format + + def _split_lines(self, text, width): + return text.splitlines() + + +class ACArgumentParser(argparse.ArgumentParser): + """Custom argparse class to override error method to change default help text.""" + + def __init__(self, + prog=None, + usage=None, + description=None, + epilog=None, + parents=None, + formatter_class=ACHelpFormatter, + prefix_chars='-', + fromfile_prefix_chars=None, + argument_default=None, + conflict_handler='error', + add_help=True, + allow_abbrev=True): + + super().__init__(prog=prog, + usage=usage, + description=description, + epilog=epilog, + parents=parents, + formatter_class=formatter_class, + prefix_chars=prefix_chars, + fromfile_prefix_chars=fromfile_prefix_chars, + argument_default=argument_default, + conflict_handler=conflict_handler, + add_help=add_help, + allow_abbrev=allow_abbrev) + register_custom_actions(self) + + self._custom_error_message = '' + + def set_custom_message(self, custom_message=''): + """ + Allows an error message override to the error() function, useful when forcing a + re-parse of arguments with newly required parameters + """ + self._custom_error_message = custom_message + + def error(self, message): + """Custom error override.""" + if len(self._custom_error_message) > 0: + message = self._custom_error_message + self._custom_error_message = '' + + lines = message.split('\n') + linum = 0 + formatted_message = '' + for line in lines: + if linum == 0: + formatted_message = 'Error: ' + line + else: + formatted_message += '\n ' + line + linum += 1 + + sys.stderr.write(Fore.LIGHTRED_EX + '{}\n\n'.format(formatted_message) + Fore.RESET) + self.print_help() + sys.exit(1) + + def format_help(self): + """Copy of format_help() from argparse.ArgumentParser with tweaks to separately display required parameters""" + formatter = self._get_formatter() + + # usage + formatter.add_usage(self.usage, self._actions, + self._mutually_exclusive_groups) + + # description + formatter.add_text(self.description) + + # positionals, optionals and user-defined groups + for action_group in self._action_groups: + if action_group.title == 'optional arguments': + # check if the arguments are required, group accordingly + req_args = [] + opt_args = [] + for action in action_group._group_actions: + if action.required: + req_args.append(action) + else: + opt_args.append(action) + + # separately display required arguments + formatter.start_section('required arguments') + formatter.add_text(action_group.description) + formatter.add_arguments(req_args) + formatter.end_section() + + # now display truly optional arguments + formatter.start_section(action_group.title) + formatter.add_text(action_group.description) + formatter.add_arguments(opt_args) + formatter.end_section() + else: + formatter.start_section(action_group.title) + formatter.add_text(action_group.description) + formatter.add_arguments(action_group._group_actions) + formatter.end_section() + + # epilog + formatter.add_text(self.epilog) + + # determine help from format above + return formatter.format_help() + + def _get_nargs_pattern(self, action): + # Override _get_nargs_pattern behavior to use the nargs ranges provided by AutoCompleter + if isinstance(action, _RangeAction) and \ + action.nargs_min is not None and action.nargs_max is not None: + nargs_pattern = '(-*A{{{},{}}}-*)'.format(action.nargs_min, action.nargs_max) + + # if this is an optional action, -- is not allowed + if action.option_strings: + nargs_pattern = nargs_pattern.replace('-*', '') + nargs_pattern = nargs_pattern.replace('-', '') + return nargs_pattern + return super(ACArgumentParser, self)._get_nargs_pattern(action) + + def _match_argument(self, action, arg_strings_pattern): + # match the pattern for this action to the arg strings + nargs_pattern = self._get_nargs_pattern(action) + match = _re.match(nargs_pattern, arg_strings_pattern) + + # raise an exception if we weren't able to find a match + if match is None: + if isinstance(action, _RangeAction) and \ + action.nargs_min is not None and action.nargs_max is not None: + raise ArgumentError(action, + 'Expected between {} and {} arguments'.format(action.nargs_min, action.nargs_max)) + + return super(ACArgumentParser, self)._match_argument(action, arg_strings_pattern) diff --git a/examples/tab_autocompletion.py b/examples/tab_autocompletion.py index fab4ce01..8cbe7f0a 100755 --- a/examples/tab_autocompletion.py +++ b/examples/tab_autocompletion.py @@ -4,9 +4,10 @@ """ import argparse import AutoCompleter +from typing import List import cmd2 -from cmd2 import with_argparser +from cmd2 import with_argparser, with_category # List of strings used with flag and index based completion functions food_item_strs = ['Pizza', 'Hamburger', 'Ham', 'Potato'] @@ -16,9 +17,60 @@ sport_item_strs = ['Bat', 'Basket', 'Basketball', 'Football'] class TabCompleteExample(cmd2.Cmd): """ Example cmd2 application where we a base command which has a couple subcommands.""" + CAT_AUTOCOMPLETE = 'AutoComplete Examples' + def __init__(self): super().__init__() + # For mocking a data source for the example commands + ratings_types = ['G', 'PG', 'PG-13', 'R', 'NC-17'] + static_list_directors = ['J. J. Abrams', 'Irvin Kershner', 'George Lucas', 'Richard Marquand', + 'Rian Johnson', 'Gareth Edwards'] + actors = ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', + 'Anthony Daniels', 'Adam Driver', 'Daisy Ridley', 'John Boyega', 'Oscar Isaac', + 'Lupita Nyong\'o', 'Andy Serkis', 'Liam Neeson', 'Ewan McGregor', 'Natalie Portman', + 'Jake Lloyd', 'Hayden Christensen', 'Christopher Lee'] + USER_MOVIE_LIBRARY = ['ROGUE1', 'SW_EP04', 'SW_EP05'] + MOVIE_DATABASE_IDS = ['SW_EP01', 'SW_EP02', 'SW_EP03', 'ROGUE1', 'SW_EP04', + 'SW_EP05', 'SW_EP06', 'SW_EP07', 'SW_EP08', 'SW_EP09'] + MOVIE_DATABASE = {'SW_EP04': {'title': 'Star Wars: Episode IV - A New Hope', + 'rating': 'PG', + 'director': ['George Lucas'], + 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', + 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] + }, + 'SW_EP05': {'title': 'Star Wars: Episode V - The Empire Strikes Back', + 'rating': 'PG', + 'director': ['Irvin Kershner'], + 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', + 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] + }, + 'SW_EP06': {'title': 'Star Wars: Episode IV - A New Hope', + 'rating': 'PG', + 'director': ['Richard Marquand'], + 'actor': ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', + 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels'] + }, + 'SW_EP01': {'title': 'Star Wars: Episode I - The Phantom Menace', + 'rating': 'PG', + 'director': ['George Lucas'], + 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', 'Jake Lloyd'] + }, + 'SW_EP02': {'title': 'Star Wars: Episode II - Attack of the Clones', + 'rating': 'PG', + 'director': ['George Lucas'], + 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', + 'Hayden Christensen', 'Christopher Lee'] + }, + 'SW_EP03': {'title': 'Star Wars: Episode III - Revenge of the Sith', + 'rating': 'PG-13', + 'director': ['George Lucas'], + 'actor': ['Liam Neeson', 'Ewan McGregor', 'Natalie Portman', + 'Hayden Christensen'] + }, + + } + # This demonstrates a number of customizations of the AutoCompleter version of ArgumentParser # - The help output will separately group required vs optional flags # - The help output for arguments with multiple flags or with append=True is more concise @@ -32,8 +84,9 @@ class TabCompleteExample(cmd2.Cmd): '\tsingle value - maximum duration\n' '\t[a, b] - duration range') + @with_category(CAT_AUTOCOMPLETE) @with_argparser(suggest_parser) - def do_suggest(self, args): + def do_suggest(self, args) -> None: """Suggest command demonstrates argparse customizations See hybrid_suggest and orig_suggest to compare the help output. @@ -43,7 +96,7 @@ class TabCompleteExample(cmd2.Cmd): if not args.type: self.do_help('suggest') - def complete_suggest(self, text, line, begidx, endidx): + def complete_suggest(self, text: str, line: str, begidx: int, endidx: int) -> List[str]: """ Adds tab completion to media""" completer = AutoCompleter.AutoCompleter(TabCompleteExample.suggest_parser, 1) @@ -64,6 +117,8 @@ class TabCompleteExample(cmd2.Cmd): help='Duration constraint in minutes.\n' '\tsingle value - maximum duration\n' '\t[a, b] - duration range') + + @with_category(CAT_AUTOCOMPLETE) @with_argparser(suggest_parser_hybrid) def do_hybrid_suggest(self, args): if not args.type: @@ -89,12 +144,14 @@ class TabCompleteExample(cmd2.Cmd): help='Duration constraint in minutes.\n' '\tsingle value - maximum duration\n' '\t[a, b] - duration range') + @with_argparser(suggest_parser_orig) - def do_orig_suggest(self, args): + @with_category(CAT_AUTOCOMPLETE) + def do_orig_suggest(self, args) -> None: if not args.type: self.do_help('orig_suggest') - def complete_orig_suggest(self, text, line, begidx, endidx): + def complete_orig_suggest(self, text, line, begidx, endidx) -> List[str]: """ Adds tab completion to media""" completer = AutoCompleter.AutoCompleter(TabCompleteExample.suggest_parser_orig) @@ -103,27 +160,29 @@ class TabCompleteExample(cmd2.Cmd): return results - ################################################################################### # The media command demonstrates a completer with multiple layers of subcommands - # + # - This example uses a flat completion lookup dictionary - def query_actors(self): + def query_actors(self) -> List[str]: """Simulating a function that queries and returns a completion values""" - return ['Mark Hamill', 'Harrison Ford', 'Carrie Fisher', 'Alec Guinness', 'Peter Mayhew', 'Anthony Daniels', - 'Adam Driver', 'Daisy Ridley', 'John Boyega', 'Oscar Isaac', 'Lupita Nyong\'o', 'Andy Serkis'] + return TabCompleteExample.actors - def _do_media_movies(self, args): + def _do_media_movies(self, args) -> None: if not args.command: self.do_help('media movies') - - def _do_media_shows(self, args): + elif args.command == 'list': + for movie_id in TabCompleteExample.MOVIE_DATABASE: + movie = TabCompleteExample.MOVIE_DATABASE[movie_id] + print('{}\n-----------------------------\n{} ID: {}\nDirector: {}\nCast:\n {}\n\n' + .format(movie['title'], movie['rating'], movie_id, + ', '.join(movie['director']), + '\n '.join(movie['actor']))) + + def _do_media_shows(self, args) -> None: if not args.command: self.do_help('media shows') - # example choices list - ratings_types = ['G', 'PG', 'PG-13', 'R', 'NC-17'] - media_parser = AutoCompleter.ACArgumentParser() media_types_subparsers = media_parser.add_subparsers(title='Media Types', dest='type') @@ -152,6 +211,7 @@ class TabCompleteExample(cmd2.Cmd): shows_parser = media_types_subparsers.add_parser('shows') shows_parser.set_defaults(func=_do_media_shows) + @with_category(CAT_AUTOCOMPLETE) @with_argparser(media_parser) def do_media(self, args): """Media management command demonstrates multiple layers of subcommands being handled by AutoCompleter""" @@ -169,10 +229,9 @@ class TabCompleteExample(cmd2.Cmd): # name collisions. def complete_media(self, text, line, begidx, endidx): """ Adds tab completion to media""" - static_list_directors = ['J. J. Abrams', 'Irvin Kershner', 'George Lucas', 'Richard Marquand', - 'Rian Johnson', 'Gareth Edwards'] - choices = {'actor': self.query_actors, - 'director': static_list_directors} + choices = {'actor': self.query_actors, # function + 'director': TabCompleteExample.static_list_directors # static list + } completer = AutoCompleter.AutoCompleter(TabCompleteExample.media_parser, arg_choices=choices) tokens, _ = self.tokens_for_completion(line, begidx, endidx) @@ -180,6 +239,101 @@ class TabCompleteExample(cmd2.Cmd): return results + ################################################################################### + # The library command demonstrates a completer with multiple layers of subcommands + # with different completion results per sub-command + # - This demonstrates how to build a tree of completion lookups to pass down + # + # Only use this method if you absolutely need to as it dramatically + # increases the complexity and decreases readability. + + def _do_library_movie(self, args): + if not args.type or not args.command: + self.do_help('library movie') + + def _do_library_show(self, args): + if not args.type: + self.do_help('library show') + + def _query_movie_database(self, exclude=[]): + return list(set(TabCompleteExample.MOVIE_DATABASE_IDS).difference(set(exclude))) + + def _query_movie_user_library(self): + return TabCompleteExample.USER_MOVIE_LIBRARY + + library_parser = AutoCompleter.ACArgumentParser(prog='library') + + library_subcommands = library_parser.add_subparsers(title='Media Types', dest='type') + + library_movie_parser = library_subcommands.add_parser('movie') + library_movie_parser.set_defaults(func=_do_library_movie) + + library_movie_subcommands = library_movie_parser.add_subparsers(title='Command', dest='command') + + library_movie_add_parser = library_movie_subcommands.add_parser('add') + library_movie_add_parser.add_argument('movie_id', help='ID of movie to add', action='append') + + library_movie_remove_parser = library_movie_subcommands.add_parser('remove') + library_movie_remove_parser.add_argument('movie_id', help='ID of movie to remove', action='append') + + library_show_parser = library_subcommands.add_parser('show') + library_show_parser.set_defaults(func=_do_library_show) + + @with_category(CAT_AUTOCOMPLETE) + @with_argparser(library_parser) + def do_library(self, args): + """Media management command demonstrates multiple layers of subcommands being handled by AutoCompleter""" + func = getattr(args, 'func', None) + if func is not None: + # Call whatever subcommand function was selected + func(self, args) + else: + # No subcommand was provided, so call help + self.do_help('library') + + def complete_library(self, text, line, begidx, endidx): + + # this demonstrates the much more complicated scenario of having + # unique completion parameters per sub-command that use the same + # argument name. To do this we build a multi-layer nested tree + # of lookups far AutoCompleter to traverse. This nested tree must + # match the structure of the argparse parser + # + + movie_add_choices = {'movie_id': self._query_movie_database} + movie_remove_choices = {'movie_id': self._query_movie_user_library} + + # The library movie sub-parser group 'command' has 2 sub-parsers: + # 'add' and 'remove' + library_movie_command_params = \ + {'add': (movie_add_choices, None), + 'remove': (movie_remove_choices, None)} + + # The 'library movie' command has a sub-parser group called 'command' + library_movie_subcommand_groups = {'command': library_movie_command_params} + + # Mapping of a specific sub-parser of the 'type' group to a tuple. Each + # tuple has 2 values corresponding what's passed to the constructor + # parameters (arg_choices,subcmd_args_lookup) of the nested + # instance of AutoCompleter + library_type_params = {'movie': (None, library_movie_subcommand_groups), + 'show': (None, None)} + + # maps the a subcommand group to a dictionary mapping a specific + # sub-command to a tuple of (arg_choices, subcmd_args_lookup) + # + # In this example, 'library_parser' has a sub-parser group called 'type' + # under the type sub-parser group, there are 2 sub-parsers: 'movie', 'show' + library_subcommand_groups = {'type': library_type_params} + + completer = AutoCompleter.AutoCompleter(TabCompleteExample.library_parser, + subcmd_args_lookup=library_subcommand_groups) + + tokens, _ = self.tokens_for_completion(line, begidx, endidx) + results = completer.complete_command(tokens, text, line, begidx, endidx) + + return results + if __name__ == '__main__': app = TabCompleteExample() |