Merge pull request #2041 from spanglerco/shell-completion-option-values

davidism · web-flow · commit ef11be6e49e1 · 2022-03-28T08:08:28.000-07:00
Make shell completion prioritize option values over new options
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -50,6 +50,8 @@ Version 8.1.0
     value. :issue:`2124`
 -   ``to_info_dict`` will not fail if a ``ParamType`` doesn't define a
     ``name``. :issue:`2168`
+-   Shell completion prioritizes option values with option prefixes over
+    new options. :issue:`2040`
 
 
 Version 8.0.4
diff --git a/src/click/core.py b/src/click/core.py
@@ -292,6 +292,8 @@ def __init__(
         #: must be never propagated to another arguments.  This is used
         #: to implement nested parsing.
         self.protected_args: t.List[str] = []
+        #: the collected prefixes of the command's options.
+        self._opt_prefixes: t.Set[str] = set(parent._opt_prefixes) if parent else set()
 
         if obj is None and parent is not None:
             obj = parent.obj
@@ -1385,6 +1387,7 @@ def parse_args(self, ctx: Context, args: t.List[str]) -> t.List[str]:
             )
 
         ctx.args = args
+        ctx._opt_prefixes.update(parser._opt_prefixes)
         return args
 
     def invoke(self, ctx: Context) -> t.Any:
diff --git a/src/click/shell_completion.py b/src/click/shell_completion.py
@@ -448,17 +448,16 @@ def _is_incomplete_argument(ctx: Context, param: Parameter) -> bool:
     )
 
 
-def _start_of_option(value: str) -> bool:
+def _start_of_option(ctx: Context, value: str) -> bool:
     """Check if the value looks like the start of an option."""
     if not value:
         return False
 
     c = value[0]
-    # Allow "/" since that starts a path.
-    return not c.isalnum() and c != "/"
+    return c in ctx._opt_prefixes
 
 
-def _is_incomplete_option(args: t.List[str], param: Parameter) -> bool:
+def _is_incomplete_option(ctx: Context, args: t.List[str], param: Parameter) -> bool:
     """Determine if the given parameter is an option that needs a value.
 
     :param args: List of complete args before the incomplete value.
@@ -467,7 +466,7 @@ def _is_incomplete_option(args: t.List[str], param: Parameter) -> bool:
     if not isinstance(param, Option):
         return False
 
-    if param.is_flag:
+    if param.is_flag or param.count:
         return False
 
     last_option = None
@@ -476,7 +475,7 @@ def _is_incomplete_option(args: t.List[str], param: Parameter) -> bool:
         if index + 1 > param.nargs:
             break
 
-        if _start_of_option(arg):
+        if _start_of_option(ctx, arg):
             last_option = arg
 
     return last_option is not None and last_option in param.opts
@@ -551,23 +550,23 @@ def _resolve_incomplete(
     # split and discard the "=" to make completion easier.
     if incomplete == "=":
         incomplete = ""
-    elif "=" in incomplete and _start_of_option(incomplete):
+    elif "=" in incomplete and _start_of_option(ctx, incomplete):
         name, _, incomplete = incomplete.partition("=")
         args.append(name)
 
     # The "--" marker tells Click to stop treating values as options
     # even if they start with the option character. If it hasn't been
     # given and the incomplete arg looks like an option, the current
     # command will provide option name completions.
-    if "--" not in args and _start_of_option(incomplete):
+    if "--" not in args and _start_of_option(ctx, incomplete):
         return ctx.command, incomplete
 
     params = ctx.command.get_params(ctx)
 
     # If the last complete arg is an option name with an incomplete
     # value, the option will provide value completions.
     for param in params:
-        if _is_incomplete_option(args, param):
+        if _is_incomplete_option(ctx, args, param):
             return param, incomplete
 
     # It's not an option name or value. The first argument without a
diff --git a/tests/test_commands.py b/tests/test_commands.py
@@ -352,3 +352,62 @@ def deprecated_cmd():
 
     result = runner.invoke(deprecated_cmd)
     assert "DeprecationWarning:" in result.output
+
+
+def test_command_parse_args_collects_option_prefixes():
+    @click.command()
+    @click.option("+p", is_flag=True)
+    @click.option("!e", is_flag=True)
+    def test(p, e):
+        pass
+
+    ctx = click.Context(test)
+    test.parse_args(ctx, [])
+
+    assert ctx._opt_prefixes == {"-", "--", "+", "!"}
+
+
+def test_group_parse_args_collects_base_option_prefixes():
+    @click.group()
+    @click.option("~t", is_flag=True)
+    def group(t):
+        pass
+
+    @group.command()
+    @click.option("+p", is_flag=True)
+    def command1(p):
+        pass
+
+    @group.command()
+    @click.option("!e", is_flag=True)
+    def command2(e):
+        pass
+
+    ctx = click.Context(group)
+    group.parse_args(ctx, ["command1", "+p"])
+
+    assert ctx._opt_prefixes == {"-", "--", "~"}
+
+
+def test_group_invoke_collects_used_option_prefixes(runner):
+    opt_prefixes = set()
+
+    @click.group()
+    @click.option("~t", is_flag=True)
+    def group(t):
+        pass
+
+    @group.command()
+    @click.option("+p", is_flag=True)
+    @click.pass_context
+    def command1(ctx, p):
+        nonlocal opt_prefixes
+        opt_prefixes = ctx._opt_prefixes
+
+    @group.command()
+    @click.option("!e", is_flag=True)
+    def command2(e):
+        pass
+
+    runner.invoke(group, ["command1"])
+    assert opt_prefixes == {"-", "--", "~", "+"}
diff --git a/tests/test_context.py b/tests/test_context.py
@@ -365,3 +365,11 @@ def cli(ctx, option):
 
     rv = runner.invoke(cli, standalone_mode=False, **invoke_args)
     assert rv.return_value == expect
+
+
+def test_propagate_opt_prefixes():
+    parent = click.Context(click.Command("test"))
+    parent._opt_prefixes = {"-", "--", "!"}
+    ctx = click.Context(click.Command("test2"), parent=parent)
+
+    assert ctx._opt_prefixes == {"-", "--", "!"}
diff --git a/tests/test_parser.py b/tests/test_parser.py
@@ -1,5 +1,7 @@
 import pytest
 
+import click
+from click.parser import OptionParser
 from click.parser import split_arg_string
 
 
@@ -15,3 +17,16 @@
 )
 def test_split_arg_string(value, expect):
     assert split_arg_string(value) == expect
+
+
+def test_parser_default_prefixes():
+    parser = OptionParser()
+    assert parser._opt_prefixes == {"-", "--"}
+
+
+def test_parser_collects_prefixes():
+    ctx = click.Context(click.Command("test"))
+    parser = OptionParser(ctx)
+    click.Option("+p", is_flag=True).add_to_parser(parser, ctx)
+    click.Option("!e", is_flag=True).add_to_parser(parser, ctx)
+    assert parser._opt_prefixes == {"-", "--", "+", "!"}
diff --git a/tests/test_shell_completion.py b/tests/test_shell_completion.py
@@ -110,6 +110,45 @@ def test_type_choice():
     assert _get_words(cli, ["-c"], "a2") == ["a2"]
 
 
+def test_choice_special_characters():
+    cli = Command("cli", params=[Option(["-c"], type=Choice(["!1", "!2", "+3"]))])
+    assert _get_words(cli, ["-c"], "") == ["!1", "!2", "+3"]
+    assert _get_words(cli, ["-c"], "!") == ["!1", "!2"]
+    assert _get_words(cli, ["-c"], "!2") == ["!2"]
+
+
+def test_choice_conflicting_prefix():
+    cli = Command(
+        "cli",
+        params=[
+            Option(["-c"], type=Choice(["!1", "!2", "+3"])),
+            Option(["+p"], is_flag=True),
+        ],
+    )
+    assert _get_words(cli, ["-c"], "") == ["!1", "!2", "+3"]
+    assert _get_words(cli, ["-c"], "+") == ["+p"]
+
+
+def test_option_count():
+    cli = Command("cli", params=[Option(["-c"], count=True)])
+    assert _get_words(cli, ["-c"], "") == []
+    assert _get_words(cli, ["-c"], "-") == ["--help"]
+
+
+def test_option_optional():
+    cli = Command(
+        "cli",
+        add_help_option=False,
+        params=[
+            Option(["--name"], is_flag=False, flag_value="value"),
+            Option(["--flag"], is_flag=True),
+        ],
+    )
+    assert _get_words(cli, ["--name"], "") == []
+    assert _get_words(cli, ["--name"], "-") == ["--flag"]
+    assert _get_words(cli, ["--name", "--flag"], "-") == []
+
+
 @pytest.mark.parametrize(
     ("type", "expect"),
     [(File(), "file"), (Path(), "file"), (Path(file_okay=False), "dir")],
