Print documentation for flags

After writing a flag on the shell, pressing "tab" will print out the
documentation for that flag.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=191899137

java/google/registry/tools/ShellCommand.java

@@ -19,11 +19,11 @@ import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.nio.charset.StandardCharsets.US_ASCII;
 
 import com.beust.jcommander.JCommander;
+import com.beust.jcommander.ParameterDescription;
 import com.beust.jcommander.Parameters;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableTable;
-import com.google.common.collect.ImmutableSetMultimap;
 import java.io.BufferedReader;
 import java.io.File;
 import java.io.IOException;
@@ -33,6 +33,9 @@ import java.io.StreamTokenizer;
 import java.io.StringReader;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map.Entry;
+import java.util.Optional;
+import javax.annotation.Nullable;
 import jline.Completor;
 import jline.ConsoleReader;
 import jline.ConsoleReaderInputStream;
@@ -61,8 +64,7 @@ public class ShellCommand implements Command {
       consoleReader = new ConsoleReader();
       // There are 104 different commands. We want the threshold to be more than that
       consoleReader.setAutoprintThreshhold(200);
-      // Setting the prompt to a temporary value - will include the environment once that is set
+      consoleReader.setDefaultPrompt("nom > ");
-      consoleReader.setDefaultPrompt("nom@??? > ");
       consoleReader.setHistory(new History(new File(USER_HOME.value(), HISTORY_FILE)));
       in = new ConsoleReaderInputStream(consoleReader);
     } else {
@@ -148,33 +150,74 @@ public class ShellCommand implements Command {
   @VisibleForTesting
   static class JCommanderCompletor implements Completor {
 
-    private final ImmutableSet<String> commands;
+    /**
-    private final ImmutableSetMultimap<String, String> commandArguments;
+     * Documentation for all the known command + argument combinations.
+     *
+     * <p>For every command, has documentation for all flags and for the "main" parameters (the ones
+     * that don't have a flag)
+     *
+     * <p>The order is: row is the command, col is the flag.
+     *
+     * <p>The flag documentations are keyed to the full flag - including any dashes (so "--flag"
+     * rather than "flag").
+     *
+     * <p>The "main" parameter documentation is keyed to "". Every command has documentation for the
+     * main parameters, even if it doesn't accept any (if it doesn't accept any, the documentation
+     * is "no documentation available"). That means every command has at least one full table cell
+     * (the "" key, for the main parameter). THIS IS IMPORTANT - otherwise the command won't appear
+     * in {@link ImmutableTable#rowKeySet}.
+     */
+    private final ImmutableTable<String, String, String> commandFlagDocs;
+
     private final FileNameCompletor filenameCompletor = new FileNameCompletor();
 
+    /**
+     * Populates the completions and documentation based on the JCommander.
+     *
+     * The input data is copied, so changing the jcommander after creation of the
+     * JCommanderCompletor doesn't change the completions.
+     */
     JCommanderCompletor(JCommander jcommander) {
-      commands = ImmutableSet.copyOf(jcommander.getCommands().keySet());
+      ImmutableTable.Builder<String, String, String> builder =
-      ImmutableSetMultimap.Builder<String, String> builder = new ImmutableSetMultimap.Builder<>();
+          new ImmutableTable.Builder<>();
-      jcommander
+
-          .getCommands()
+      // Go over all the commands
-          .entrySet()
+      for (Entry<String, JCommander> entry : jcommander.getCommands().entrySet()) {
-          .forEach(
+        String command = entry.getKey();
-              entry -> {
+        JCommander subCommander = entry.getValue();
-                builder.putAll(
+
-                    entry.getKey(),
+        // Add the "main" parameters documentation
-                    entry
+        builder.put(command, "", createDocText(subCommander.getMainParameter()));
-                        .getValue()
+
-                        .getParameters()
+        // For each command - go over the parameters (arguments / flags)
-                        .stream()
+        for (ParameterDescription parameter : subCommander.getParameters()) {
-                        .flatMap(p -> Arrays.stream(p.getParameter().names()))
+          String documentation = createDocText(parameter);
-                        .collect(toImmutableList()));
+
-              });
+          // For each parameter - go over all the "flag" names of that parameter (e.g., -o and
-      commandArguments = builder.build();
+          // --output being aliases of the same parameter) and populate each one
+          Arrays.stream(parameter.getParameter().names())
+              .forEach(flag -> builder.put(command, flag, documentation));
+        }
+      }
+      commandFlagDocs = builder.build();
+    }
+
+    private static String createDocText(@Nullable ParameterDescription parameter) {
+      if (parameter == null) {
+        return "[None]";
+      }
+      String type = parameter.getParameterized().getGenericType().toString();
+      if (type.startsWith("class ")) {
+        type = type.substring(6);
+      }
+      return String.format("%s\n  (%s)", parameter.getDescription(), type);
     }
 
     @Override
     @SuppressWarnings({"unchecked", "rawtypes"})
     public int complete(String buffer, int location, List completions) {
+      // We just defer to the other function because of the warnings (the use of a naked List by
+      // jline)
       return completeInternal(buffer, location, completions);
     }
 
@@ -194,55 +237,118 @@ public class ShellCommand implements Command {
       if (argumentIndex < 0 || !truncatedBuffer.endsWith(parsedBuffer[argumentIndex])) {
         argumentIndex += 1;
       }
-      String argument = argumentIndex < parsedBuffer.length ? parsedBuffer[argumentIndex] : "";
+      // The argument we want to complete (only partially written, might even be empty)
-      int argumentStart = location - argument.length();
+      String partialArgument =
+          argumentIndex < parsedBuffer.length ? parsedBuffer[argumentIndex] : "";
+      int argumentStart = location - partialArgument.length();
+      // The command name. Null if we're at the first argument
+      String command = argumentIndex == 0 ? null : parsedBuffer[0];
+      // The previous argument before it - used for context. Null if we're at the first argument
+      String previousArgument = argumentIndex <= 1 ? null : parsedBuffer[argumentIndex - 1];
+
+      // If it's obviously a file path (starts with something "file path like") - complete as a file
+      if (partialArgument.startsWith("./")
+          || partialArgument.startsWith("~/")
+          || partialArgument.startsWith("/")) {
+        int offset =
+            filenameCompletor.complete(partialArgument, partialArgument.length(), completions);
+        if (offset >= 0) {
+          return argumentStart + offset;
+        }
+        return -1;
+      }
+
+      // Complete based on flag data
+      completions.addAll(getCompletions(command, previousArgument, partialArgument));
+      return argumentStart;
+    }
+
+    /**
+     * Completes a (partial) word based on the command and context.
+     *
+     * @param command the name of the command we're running. Null if not yet known (it is in 'word')
+     * @param context the previous argument for context. Null if we're the first.
+     * @param word the (partial) word to complete. Can be the command, if "command" is null, or any
+     *     "regular" argument, if "command" isn't null.
+     * @return list of all possible completions to 'word'
+     */
+    private List<String> getCompletions(
+        @Nullable String command, @Nullable String context, String word) {
 
       // Complete the first argument based on the jcommander commands
-      if (argumentIndex == 0) {
+      if (command == null) {
-        completions.addAll(getCommandCompletions(argument));
+        return getCommandCompletions(word);
-        return argumentStart;
       }
-      String commandName = parsedBuffer[0];
 
       // For the "help" command, complete the second argument based on the jcommander commands, and
       // the rest of the arguments fail to complete
-      if (commandName.equals("help")) {
+      if (command.equals("help")) {
-        if (argumentIndex >= 2) {
+        // "help" only has completion for the first argument
-          return argumentStart;
+        if (context != null) {
+          return ImmutableList.of();
         }
-        completions.addAll(getCommandCompletions(argument));
+        return getCommandCompletions(word);
-        return argumentStart;
+      }
+
+      // 'tab' on empty will show the documentation - either for the "current flag" or for the main
+      // parameters, depending on the context (the "context" being the previous argument)
+      if (word.isEmpty()) {
+        return getParameterDocCompletions(command, context, word);
       }
 
       // For existing commands, complete based on the command arguments
-      if (argument.isEmpty() || argument.startsWith("-")) {
+      if (word.startsWith("-")) {
-        completions.addAll(getArgumentCompletions(commandName, argument));
+        return getFlagCompletions(command, word);
-        return argumentStart;
       }
 
-      // However, if it's obviously not an argument (starts with something that isn't "-"), default
+      // We don't know how to complete based on context... :( So that's the best we can do
-      // to a filename.
+      return ImmutableList.of();
-      int offset = filenameCompletor.complete(argument, argument.length(), completions);
-      if (offset < 0) {
-        return argumentStart;
-      }
-      return argumentStart + offset;
     }
 
     private List<String> getCommandCompletions(String word) {
-      return commands
+      return commandFlagDocs
+          .rowKeySet()
           .stream()
           .filter(s -> s.startsWith(word))
           .map(s -> s + " ")
           .collect(toImmutableList());
     }
 
-    private List<String> getArgumentCompletions(String command, String word) {
+    private List<String> getFlagCompletions(String command, String word) {
-      return commandArguments.get(command)
+      return commandFlagDocs
+          .row(command)
+          .keySet()
           .stream()
           .filter(s -> s.startsWith(word))
           .map(s -> s + " ")
           .collect(toImmutableList());
     }
+
+    private List<String> getParameterDocCompletions(
+        String command, @Nullable String argument, String word) {
+      if (!word.isEmpty()) {
+        return ImmutableList.of();
+      }
+      return ImmutableList.of("", getParameterDoc(command, argument));
+    }
+
+    private String getParameterDoc(String command, @Nullable String previousArgument) {
+      // First, check if we want the documentation for a specific flag, or for the "main"
+      // parameters.
+      //
+      // We want documentation for a flag if the previous argument was a flag, but the value of the
+      // flag wasn't set. So if the previous argument is "--flag" then we want documentation of that
+      // flag, but if it's "--flag=value" then that flag is set and we want documentation of the
+      // main parameters.
+      boolean isFlagParameter =
+          previousArgument != null
+              && previousArgument.startsWith("-")
+              && previousArgument.indexOf('=') == -1;
+      return (isFlagParameter ? "Flag documentation: " : "Main parameter: ")
+          + Optional.ofNullable(
+                  commandFlagDocs.get(command, isFlagParameter ? previousArgument : ""))
+              .orElse("[No documentation available]")
+          + "\n";
+    }
   }
 }

javatests/google/registry/tools/ShellCommandTest.java

@@ -134,9 +134,29 @@ public class ShellCommandTest {
     performJCommanderCompletorTest("help testCommand ", 0);
   }
 
+  @Test
+  public void testCompletion_documentation() throws Exception {
+    performJCommanderCompletorTest(
+        "testCommand ",
+        0,
+        "",
+        "Main parameter: normal argument\n  (java.util.List<java.lang.String>)\n");
+    performJCommanderCompletorTest("testAnotherCommand ", 0, "", "Main parameter: [None]\n");
+    performJCommanderCompletorTest(
+        "testCommand -x ", 0, "", "Flag documentation: test parameter\n  (java.lang.String)\n");
+    performJCommanderCompletorTest(
+        "testAnotherCommand -x ", 0, "", "Flag documentation: [No documentation available]\n");
+    performJCommanderCompletorTest(
+        "testCommand x ",
+        0,
+        "",
+        "Main parameter: normal argument\n  (java.util.List<java.lang.String>)\n");
+    performJCommanderCompletorTest("testAnotherCommand x ", 0, "", "Main parameter: [None]\n");
+  }
+
   @Test
   public void testCompletion_arguments() throws Exception {
-    performJCommanderCompletorTest("testCommand ", 0, "-x ", "--xparam ", "--xorg ");
+    performJCommanderCompletorTest("testCommand -", 1, "-x ", "--xparam ", "--xorg ");
     performJCommanderCompletorTest("testCommand --wrong", 7);
     performJCommanderCompletorTest("testCommand noise  --", 2, "--xparam ", "--xorg ");
     performJCommanderCompletorTest("testAnotherCommand --o", 3);