add pdctl interactive mode autocomplete

tools/pd-ctl/pdctl/ctl.go

@@ -17,6 +17,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/chzyer/readline"
 	"github.com/mattn/go-shellwords"
@@ -37,9 +38,10 @@ type CommandFlags struct {
 var (
 	commandFlags = CommandFlags{}
 
-	detach   bool
-	interact bool
-	version  bool
+	detach            bool
+	interact          bool
+	version           bool
+	readlinecompleter *readline.PrefixCompleter
 )
 
 func init() {
@@ -117,6 +119,7 @@ func getMainCmd(args []string) *cobra.Command {
 	rootCmd.ParseFlags(args)
 	rootCmd.SetOutput(os.Stdout)
 
+	readlinecompleter = readline.NewPrefixCompleter(genCompleter(rootCmd)...)
 	return rootCmd
 }
 
@@ -156,6 +159,7 @@ func loop() {
 	l, err := readline.NewEx(&readline.Config{
 		Prompt:            "\033[31m»\033[0m ",
 		HistoryFile:       "/tmp/readline.tmp",
+		AutoComplete:      readlinecompleter,
 		InterruptPrompt:   "^C",
 		EOFPrompt:         "^D",
 		HistorySearchFold: true,
@@ -190,3 +194,23 @@ func loop() {
 		Start(args)
 	}
 }
+
+func genCompleter(cmd *cobra.Command) []readline.PrefixCompleterInterface {
+	pc := []readline.PrefixCompleterInterface{}
+	if len(cmd.Commands()) != 0 {
+		for _, v := range cmd.Commands() {
+			if v.HasFlags() {
+				flagspc := []readline.PrefixCompleterInterface{}
+				flaguseages := strings.Split(strings.Trim(v.Flags().FlagUsages(), " "), "\n")
+				for i := 0; i < len(flaguseages)-1; i++ {
+					flagspc = append(flagspc, readline.PcItem(strings.Split(strings.Trim(flaguseages[i], " "), " ")[0]))
+				}
+				flagspc = append(flagspc, genCompleter(v)...)
+				pc = append(pc, readline.PcItem(strings.Split(v.Use, " ")[0], flagspc...))
+			} else {
+				pc = append(pc, readline.PcItem(strings.Split(v.Use, " ")[0], genCompleter(v)...))
+			}
+		}
+	}
+	return pc
+}

tools/pd-ctl/pdctl/ctl_test.go

@@ -0,0 +1,87 @@
+// Copyright 2020 TiKV Project Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pdctl
+
+import (
+	"testing"
+
+	"github.com/spf13/cobra"
+)
+
+func newACommand() *cobra.Command {
+	a := &cobra.Command{
+		Use:   "testa",
+		Short: "test A command",
+	}
+	return a
+}
+
+func newBCommand() *cobra.Command {
+	a := &cobra.Command{
+		Use:   "testb",
+		Short: "test b command",
+	}
+	return a
+}
+func newDEFCommand() *cobra.Command {
+	a := &cobra.Command{
+		Use:   "testdef",
+		Short: "test def command",
+	}
+	return a
+}
+
+func newCCommand() *cobra.Command {
+	a := &cobra.Command{
+		Use:   "testc",
+		Short: "test c command",
+	}
+	return a
+}
+
+func Test_genCompleter(t *testing.T) {
+	var subcommand = []string{"testa", "testb", "testc", "testdef"}
+
+	rootCmd := &cobra.Command{
+		Use:   "roottest",
+		Short: "test root cmd",
+	}
+	rootCmd.AddCommand(newACommand(), newBCommand(), newCCommand(), newDEFCommand())
+
+	pc := genCompleter(rootCmd)
+
+	for _, cmd := range subcommand {
+		runarray := []rune(cmd)
+		inprefixarray := true
+		for _, v := range pc {
+			inprefixarray = true
+			if len(runarray) != len(v.GetName())-1 {
+				continue
+			}
+			for i := 0; i < len(runarray); i++ {
+				if runarray[i] != v.GetName()[i] {
+					inprefixarray = false
+				}
+			}
+			if inprefixarray == true {
+				break
+			}
+		}
+
+		if inprefixarray == false {
+			t.Errorf("%s not in prefix array", cmd)
+		}
+	}
+
+}