Accept the first character as shortcut on interruption prompts

On interruption, this makes buildstream to accept the first character of all
the possible choices as if it was the full command. This behavior has been
added to the failure screen and to the interruption screen.

Fixes https://gitlab.com/BuildStream/buildstream/issues/130

2 files changed, 74 insertions, 12 deletions

diff --git a/buildstream/_frontend/main.py b/buildstream/_frontend/main.py
index 63051af66..70ff24fa8 100644
--- a/buildstream/_frontend/main.py
+++ b/buildstream/_frontend/main.py
@@ -23,6 +23,7 @@ import click
 import pkg_resources  # From setuptools
 from contextlib import contextmanager
 from blessings import Terminal
+from click import UsageError
 
 # Import buildstream public symbols
 from .. import Scope, Consistency
@@ -891,16 +892,16 @@ class App():
             click.echo("\nUser interrupted with ^C\n" +
                        "\n"
                        "Choose one of the following options:\n" +
-                       "  continue  - Continue queueing jobs as much as possible\n" +
-                       "  quit      - Exit after all ongoing jobs complete\n" +
-                       "  terminate - Terminate any ongoing jobs and exit\n" +
+                       "  (c)ontinue  - Continue queueing jobs as much as possible\n" +
+                       "  (q)uit      - Exit after all ongoing jobs complete\n" +
+                       "  (t)erminate - Terminate any ongoing jobs and exit\n" +
                        "\n" +
                        "Pressing ^C again will terminate jobs and exit\n",
                        err=True)
 
             try:
                 choice = click.prompt("Choice:",
-                                      type=click.Choice(['continue', 'quit', 'terminate']),
+                                      value_proc=prefix_choice_value_proc(['continue', 'quit', 'terminate']),
                                       default='continue', err=True)
             except click.Abort:
                 # Ensure a newline after automatically printed '^C'
@@ -961,14 +962,14 @@ class App():
             summary = ("\n{} failure on element: {}\n".format(failure.action_name, element.name) +
                        "\n" +
                        "Choose one of the following options:\n" +
-                       "  continue  - Continue queueing jobs as much as possible\n" +
-                       "  quit      - Exit after all ongoing jobs complete\n" +
-                       "  terminate - Terminate any ongoing jobs and exit\n" +
-                       "  retry     - Retry this job\n")
+                       "  (c)ontinue  - Continue queueing jobs as much as possible\n" +
+                       "  (q)uit      - Exit after all ongoing jobs complete\n" +
+                       "  (t)erminate - Terminate any ongoing jobs and exit\n" +
+                       "  (r)etry     - Retry this job\n")
             if failure.logfile:
-                summary += "  log       - View the full log file\n"
+                summary += "  (l)og       - View the full log file\n"
             if failure.sandbox:
-                summary += "  shell     - Drop into a shell in the failed build sandbox\n"
+                summary += "  (s)hell     - Drop into a shell in the failed build sandbox\n"
             summary += "\nPressing ^C will terminate jobs and exit\n"
 
             choices = ['continue', 'quit', 'terminate', 'retry']
@@ -982,8 +983,8 @@ class App():
                 click.echo(summary, err=True)
 
                 try:
-                    choice = click.prompt("Choice:", type=click.Choice(choices),
-                                          default='continue', err=True)
+                    choice = click.prompt("Choice:", default='continue', err=True,
+                                          value_proc=prefix_choice_value_proc(choices))
                 except click.Abort:
                     # Ensure a newline after automatically printed '^C'
                     click.echo("", err=True)
@@ -1152,3 +1153,30 @@ class App():
         self.maybe_render_status()
         self.scheduler.resume_jobs()
         self.scheduler.connect_signals()
+
+
+#
+# Return a value processor for partial choice matching.
+# The returned values processor will test the passed value with all the item
+# in the 'choices' list. If the value is a prefix of one of the 'choices'
+# element, the element is returned. If no element or several elements match
+# the same input, a 'click.UsageError' exception is raised with a description
+# of the error.
+#
+# Note that Click expect user input errors to be signaled by raising a
+# 'click.UsageError' exception. That way, Click display an error message and
+# ask for a new input.
+#
+def prefix_choice_value_proc(choices):
+
+    def value_proc(user_input):
+        remaining_candidate = [choice for choice in choices if choice.startswith(user_input)]
+
+        if len(remaining_candidate) == 0:
+            raise UsageError("Expected one of {}, got {}".format(choices, user_input))
+        elif len(remaining_candidate) == 1:
+            return remaining_candidate[0]
+        else:
+            raise UsageError("Ambiguous input. '{}' can refer to one of {}".format(user_input, remaining_candidate))
+
+    return value_proc
diff --git a/tests/frontend/main.py b/tests/frontend/main.py
new file mode 100644
index 000000000..9ba552c60
--- /dev/null
+++ b/tests/frontend/main.py
@@ -0,0 +1,34 @@
+from buildstream._frontend.main import prefix_choice_value_proc
+
+import pytest
+import click
+
+
+def test_prefix_choice_value_proc_full_match():
+    value_proc = prefix_choice_value_proc(['foo', 'bar', 'baz'])
+
+    assert("foo" == value_proc("foo"))
+    assert("bar" == value_proc("bar"))
+    assert("baz" == value_proc("baz"))
+
+
+def test_prefix_choice_value_proc_prefix_match():
+    value_proc = prefix_choice_value_proc(['foo'])
+
+    assert ("foo" == value_proc("f"))
+
+
+def test_prefix_choice_value_proc_ambigous_match():
+    value_proc = prefix_choice_value_proc(['bar', 'baz'])
+
+    assert ("bar" == value_proc("bar"))
+    assert ("baz" == value_proc("baz"))
+    with pytest.raises(click.UsageError):
+        value_proc("ba")
+
+
+def test_prefix_choice_value_proc_value_not_in_choices():
+    value_proc = prefix_choice_value_proc(['bar', 'baz'])
+
+    with pytest.raises(click.UsageError):
+        value_proc("foo")