Issue #535: Bash completion for click.Choice()

Added support for bash compeletion of
type=click.Choice Options and Arguments.

Added comprehensive tests for bash completion.
There are many scenarios present for arguments.

.gitignore

@@ -8,3 +8,5 @@ build
 docs/_build
 click.egg-info
 .tox
+.cache
+.idea

click/_bashcomplete.py

@@ -1,9 +1,15 @@
+import collections
+import copy
 import os
 import re
+
 from .utils import echo
 from .parser import split_arg_string
-from .core import MultiCommand, Option
+from .core import MultiCommand, Option, Argument
+from .types import Choice
+
 
+WORDBREAK = '='
 
 COMPLETION_SCRIPT = '''
 %(complete_func)s() {
@@ -29,6 +35,13 @@ def get_completion_script(prog_name, complete_var):
 
 
 def resolve_ctx(cli, prog_name, args):
+    """
+    Parse into a hierarchy of contexts. Contexts are connected through the parent variable.
+    :param cli: command definition
+    :param prog_name: the program that is running
+    :param args: full list of args
+    :return: the final context/command parsed
+    """
     ctx = cli.make_context(prog_name, args, resilient_parsing=True)
     while ctx.protected_args + ctx.args and isinstance(ctx.command, MultiCommand):
         a = ctx.protected_args + ctx.args
@@ -39,19 +52,103 @@ def resolve_ctx(cli, prog_name, args):
     return ctx
 
 
+def start_of_option(param_str):
+    """
+    :param param_str: param_str to check
+    :return: whether or not this is the start of an option declaration (i.e. starts "-" or "--")
+    """
+    return param_str and param_str[:1] == '-'
+
+
+def is_incomplete_option(all_args, cmd_param):
+    """
+    :param all_args: the full original list of args supplied
+    :param cmd_param: the current command paramter
+    :return: whether or not the last option declaration (i.e. starts "-" or "--") is incomplete and
+    corresponds to this cmd_param. In other words whether this cmd_param option can still accept
+    values
+    """
+    if cmd_param.is_flag:
+        return False
+    last_option = None
+    for index, arg_str in enumerate(reversed([arg for arg in all_args if arg != WORDBREAK])):
+        if index + 1 > cmd_param.nargs:
+            break
+        if start_of_option(arg_str):
+            last_option = arg_str
+
+    return True if last_option and last_option in cmd_param.opts else False
+
+
+def is_incomplete_argument(current_params, cmd_param):
+    """
+    :param current_params: the current params and values for this argument as already entered
+    :param cmd_param: the current command parameter
+    :return: whether or not the last argument is incomplete and corresponds to this cmd_param. In
+    other words whether or not the this cmd_param argument can still accept values
+    """
+    current_param_values = current_params[cmd_param.name]
+    if current_param_values is None:
+        return True
+    if cmd_param.nargs == -1:
+        return True
+    if isinstance(current_param_values, collections.Iterable) \
+            and cmd_param.nargs > 1 and len(current_param_values) < cmd_param.nargs:
+        return True
+    return False
+
+
 def get_choices(cli, prog_name, args, incomplete):
+    """
+    :param cli: command definition
+    :param prog_name: the program that is running
+    :param args: full list of args
+    :param incomplete: the incomplete text to autocomplete
+    :return: all the possible completions for the incomplete
+    """
+    all_args = copy.deepcopy(args)
+
     ctx = resolve_ctx(cli, prog_name, args)
     if ctx is None:
         return
 
+    # In newer versions of bash long opts with '='s are partitioned, but it's easier to parse
+    # without the '='
+    if start_of_option(incomplete) and WORDBREAK in incomplete:
+        partition_incomplete = incomplete.partition(WORDBREAK)
+        all_args.append(partition_incomplete[0])
+        incomplete = partition_incomplete[2]
+    elif incomplete == WORDBREAK:
+        incomplete = ''
+
     choices = []
-    if incomplete and not incomplete[:1].isalnum():
+    found_param = False
+    if start_of_option(incomplete):
+        # completions for options
         for param in ctx.command.params:
-            if not isinstance(param, Option):
-                continue
-            choices.extend(param.opts)
-            choices.extend(param.secondary_opts)
-    elif isinstance(ctx.command, MultiCommand):
+            if isinstance(param, Option):
+                choices.extend([param_opt for param_opt in param.opts + param.secondary_opts
+                                if param_opt not in all_args or param.multiple])
+        found_param = True
+    if not found_param:
+        # completion for option values by choices
+        for cmd_param in ctx.command.params:
+            if isinstance(cmd_param, Option) and is_incomplete_option(all_args, cmd_param):
+                if isinstance(cmd_param.type, Choice):
+                    choices.extend(cmd_param.type.choices)
+                found_param = True
+                break
+    if not found_param:
+        # completion for argument values by choices
+        for cmd_param in ctx.command.params:
+            if isinstance(cmd_param, Argument) and is_incomplete_argument(ctx.params, cmd_param):
+                if isinstance(cmd_param.type, Choice):
+                    choices.extend(cmd_param.type.choices)
+                found_param = True
+                break
+
+    if not found_param and isinstance(ctx.command, MultiCommand):
+        # completion for any subcommands
         choices.extend(ctx.command.list_commands(ctx))
 
     for item in choices:

docs/bashcomplete.rst

@@ -13,16 +13,18 @@ Limitations
 Bash completion is only available if a script has been installed properly,
 and not executed through the ``python`` command.  For information about
 how to do that, see :ref:`setuptools-integration`.  Also, Click currently
-only supports completion for Bash.
+only supports completion for Bash. Zsh support is available through Zsh's
+bash completion compatibility mode.
 
 Currently, Bash completion is an internal feature that is not customizable.
 This might be relaxed in future versions.
 
 What it Completes
 -----------------
 
-Generally, the Bash completion support will complete subcommands and
-parameters.  Subcommands are always listed whereas parameters only if at
+Generally, the Bash completion support will complete subcommands, options
+as well as any option and argument values where the type is click.Choice.
+Subcommands and choices are always listed whereas parameters only if at
 least a dash has been provided.  Example::
 
     $ repo <TAB><TAB>
@@ -67,3 +69,14 @@ This can be easily accomplished::
 And then you would put this into your bashrc instead::
 
     . /path/to/foo-bar-complete.sh
+
+Zsh Compatibility
+----------------
+
+To enable Bash completion in Zsh, add the following lines to your .zshrc:
+
+    autoload bashcompinit
+    bashcompinit
+
+See https://github.com/pallets/click/issues/323 for more information on
+this issue.

tests/test_bashcomplete.py

@@ -14,6 +14,42 @@ def cli(local_opt):
     assert list(get_choices(cli, 'lol', [], '')) == []
 
 
+def test_boolean_flag():
+    @click.command()
+    @click.option('--shout/--no-shout', default=False)
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '-')) == ['--shout', '--no-shout']
+
+
+def test_multi_value_option():
+    @click.group()
+    @click.option('--pos', nargs=2, type=float)
+    def cli(local_opt):
+        pass
+
+    @cli.command()
+    @click.option('--local-opt')
+    def sub(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '-')) == ['--pos']
+    assert list(get_choices(cli, 'lol', ['--pos'], '')) == []
+    assert list(get_choices(cli, 'lol', ['--pos', '1.0'], '')) == []
+    assert list(get_choices(cli, 'lol', ['--pos', '1.0', '1.0'], '')) == ['sub']
+
+
+def test_multi_option():
+    @click.command()
+    @click.option('--message', '-m', multiple=True)
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '-')) == ['--message', '-m']
+    assert list(get_choices(cli, 'lol', ['-m'], '')) == []
+
+
 def test_small_chain():
     @click.group()
     @click.option('--global-opt')
@@ -60,3 +96,134 @@ def csub(csub_opt):
     assert list(get_choices(cli, 'lol', ['asub', 'bsub'], '')) == ['csub']
     assert list(get_choices(cli, 'lol', ['asub', 'bsub', 'csub'], '-')) == ['--csub-opt']
     assert list(get_choices(cli, 'lol', ['asub', 'bsub', 'csub'], '')) == []
+
+
+def test_argument_choice():
+    @click.command()
+    @click.argument('arg1', required=False, type=click.Choice(['arg11', 'arg12']))
+    @click.argument('arg2', required=False, type=click.Choice(['arg21', 'arg22']))
+    @click.argument('arg3', required=False, type=click.Choice(['arg', 'argument']))
+    def cli():
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '')) == ['arg11', 'arg12']
+    assert list(get_choices(cli, 'lol', [], 'arg')) == ['arg11', 'arg12']
+    assert list(get_choices(cli, 'lol', ['arg11'], '')) == ['arg21', 'arg22']
+    assert list(get_choices(cli, 'lol', ['arg12', 'arg21'], '')) == ['arg', 'argument']
+    assert list(get_choices(cli, 'lol', ['arg12', 'arg21'], 'argu')) == ['argument']
+
+
+def test_option_choice():
+    @click.command()
+    @click.option('--opt1', type=click.Choice(['opt11', 'opt12']))
+    @click.option('--opt2', type=click.Choice(['opt21', 'opt22']))
+    @click.option('--opt3', type=click.Choice(['opt', 'option']))
+    def cli():
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '-')) == ['--opt1', '--opt2', '--opt3']
+    assert list(get_choices(cli, 'lol', [], '--opt')) == ['--opt1', '--opt2', '--opt3']
+    assert list(get_choices(cli, 'lol', [], '--opt1=')) == ['opt11', 'opt12']
+    assert list(get_choices(cli, 'lol', [], '--opt2=')) == ['opt21', 'opt22']
+    assert list(get_choices(cli, 'lol', ['--opt2'], '=')) == ['opt21', 'opt22']
+    assert list(get_choices(cli, 'lol', ['--opt2', '='], 'opt')) == ['opt21', 'opt22']
+    assert list(get_choices(cli, 'lol', ['--opt1'], '')) == ['opt11', 'opt12']
+    assert list(get_choices(cli, 'lol', ['--opt2'], '')) == ['opt21', 'opt22']
+    assert list(get_choices(cli, 'lol', ['--opt1', 'opt11', '--opt2'], '')) == ['opt21', 'opt22']
+    assert list(get_choices(cli, 'lol', ['--opt2', 'opt21'], '-')) == ['--opt1', '--opt3']
+    assert list(get_choices(cli, 'lol', ['--opt1', 'opt11'], '-')) == ['--opt2', '--opt3']
+    assert list(get_choices(cli, 'lol', ['--opt1'], 'opt')) == ['opt11', 'opt12']
+    assert list(get_choices(cli, 'lol', ['--opt3'], 'opti')) == ['option']
+
+    assert list(get_choices(cli, 'lol', ['--opt1', 'invalid_opt'], '-')) == ['--opt2', '--opt3']
+
+
+def test_option_and_arg_choice():
+    @click.command()
+    @click.option('--opt1', type=click.Choice(['opt11', 'opt12']))
+    @click.argument('arg1', required=False, type=click.Choice(['arg11', 'arg12']))
+    @click.option('--opt2', type=click.Choice(['opt21', 'opt22']))
+    def cli():
+        pass
+
+    assert list(get_choices(cli, 'lol', ['--opt1'], '')) == ['opt11', 'opt12']
+    assert list(get_choices(cli, 'lol', [''], '--opt1=')) == ['opt11', 'opt12']
+    assert list(get_choices(cli, 'lol', [], '')) == ['arg11', 'arg12']
+    assert list(get_choices(cli, 'lol', ['--opt2'], '')) == ['opt21', 'opt22']
+
+
+def test_boolean_flag_choice():
+    @click.command()
+    @click.option('--shout/--no-shout', default=False)
+    @click.argument('arg', required=False, type=click.Choice(['arg1', 'arg2']))
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', [], '-')) == ['--shout', '--no-shout']
+    assert list(get_choices(cli, 'lol', ['--shout'], '')) == ['arg1', 'arg2']
+
+
+def test_multi_value_option_choice():
+    @click.command()
+    @click.option('--pos', nargs=2, type=click.Choice(['pos1', 'pos2']))
+    @click.argument('arg', required=False, type=click.Choice(['arg1', 'arg2']))
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', ['--pos'], '')) == ['pos1', 'pos2']
+    assert list(get_choices(cli, 'lol', ['--pos', 'pos1'], '')) == ['pos1', 'pos2']
+    assert list(get_choices(cli, 'lol', ['--pos', 'pos1', 'pos2'], '')) == ['arg1', 'arg2']
+    assert list(get_choices(cli, 'lol', ['--pos', 'pos1', 'pos2', 'arg1'], '')) == []
+
+
+def test_multi_option_choice():
+    @click.command()
+    @click.option('--message', '-m', multiple=True, type=click.Choice(['m1', 'm2']))
+    @click.argument('arg', required=False, type=click.Choice(['arg1', 'arg2']))
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', ['-m'], '')) == ['m1', 'm2']
+    assert list(get_choices(cli, 'lol', ['-m', 'm1', '-m'], '')) == ['m1', 'm2']
+    assert list(get_choices(cli, 'lol', ['-m', 'm1'], '')) == ['arg1', 'arg2']
+
+
+def test_variadic_argument_choice():
+    @click.command()
+    @click.argument('src', nargs=-1, type=click.Choice(['src1', 'src2']))
+    def cli(local_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', ['src1', 'src2'], '')) == ['src1', 'src2']
+
+
+def test_long_chain_choice():
+    @click.group()
+    def cli():
+        pass
+
+    @cli.group('sub')
+    @click.option('--sub-opt', type=click.Choice(['subopt1', 'subopt2']))
+    @click.argument('sub-arg', required=False, type=click.Choice(['subarg1', 'subarg2']))
+    def sub(sub_opt):
+        pass
+
+    @sub.command('bsub')
+    @click.option('--bsub-opt', type=click.Choice(['bsubopt1', 'bsubopt2']))
+    @click.argument('bsub-arg1', required=False, type=click.Choice(['bsubarg1', 'bsubarg2']))
+    @click.argument('bbsub-arg2', required=False, type=click.Choice(['bbsubarg1', 'bbsubarg2']))
+    def bsub(bsub_opt):
+        pass
+
+    assert list(get_choices(cli, 'lol', ['sub'], '')) == ['subarg1', 'subarg2']
+    assert list(get_choices(cli, 'lol', ['sub', '--sub-opt'], '')) == ['subopt1', 'subopt2']
+    assert list(get_choices(cli, 'lol', ['sub', '--sub-opt', 'subopt1'], '')) == \
+        ['subarg1', 'subarg2']
+    assert list(get_choices(cli, 'lol',
+        ['sub', '--sub-opt', 'subopt1', 'subarg1', 'bsub'], '-')) == ['--bsub-opt']
+    assert list(get_choices(cli, 'lol',
+        ['sub', '--sub-opt', 'subopt1', 'subarg1', 'bsub', '--bsub-opt'], '')) == \
+        ['bsubopt1', 'bsubopt2']
+    assert list(get_choices(cli, 'lol',
+        ['sub', '--sub-opt', 'subopt1', 'subarg1', 'bsub', '--bsub-opt', 'bsubopt1', 'bsubarg1'],
+                            '')) == ['bbsubarg1', 'bbsubarg2']