diff --git a/tools/pd-ctl/pdctl/ctl.go b/tools/pd-ctl/pdctl/ctl.go index 52278d527bd..607b9d04786 100644 --- a/tools/pd-ctl/pdctl/ctl.go +++ b/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,22 @@ func loop() { Start(args) } } + +func genCompleter(cmd *cobra.Command) []readline.PrefixCompleterInterface { + pc := []readline.PrefixCompleterInterface{} + + for _, v := range cmd.Commands() { + if v.HasFlags() { + flagsPc := []readline.PrefixCompleterInterface{} + flagUsages := strings.Split(strings.Trim(v.Flags().FlagUsages(), " "), "\n") + for i := 0; i < len(flagUsages)-1; i++ { + flagsPc = append(flagsPc, readline.PcItem(strings.Split(strings.Trim(flagUsages[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 +} diff --git a/tools/pd-ctl/pdctl/ctl_test.go b/tools/pd-ctl/pdctl/ctl_test.go new file mode 100644 index 00000000000..caf6b8a7aa0 --- /dev/null +++ b/tools/pd-ctl/pdctl/ctl_test.go @@ -0,0 +1,70 @@ +// 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 newCommand(usage, short string) *cobra.Command { + cmd := &cobra.Command{ + Use: usage, + Short: short, + } + return cmd +} + +func TestGenCompleter(t *testing.T) { + var subCommand = []string{"testa", "testb", "testc", "testdef"} + + rootCmd := &cobra.Command{ + Use: "roottest", + Short: "test root cmd", + } + + cmdA := newCommand("testa", "test a command") + cmdB := newCommand("testb", "test b command") + cmdC := newCommand("testc", "test c command") + cmdDEF := newCommand("testdef", "test def command") + + rootCmd.AddCommand(cmdA, cmdB, cmdC, cmdDEF) + + 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) + } + } + +}