Repository: stern/stern
Pull request #136 (2021-06-20 14:29:59)

Move custiom completions implemented with bash to Go
With this PR, all custom completions implemented with bash are removed, and new custom completions implemented with Go are added. In addition, stern now supports dynamic completion for fish shell.

Closes https://github.com/stern/stern/issues/20.

superbrothers: ping @floryut 

Patch:
diff --git a/README.md b/README.md
index 165655b..6db3146 100644
--- a/README.md
+++ b/README.md
@@ -58,7 +58,7 @@ The `pod` query is a regular expression so you could provide `"web-\w"` to tail
 -----------------------------|-----------|---------
  `--all-namespaces`, `-A`    | `false`   | If present, tail across all namespaces. A specific namespace is ignored even if specified with --namespace.
  `--color`                   | `auto`    | Force set color output. 'auto':  colorize if tty attached, 'always': always colorize, 'never': never colorize.
- `--completion`              |           | Output stern command-line completion code for the specified shell. Can be 'bash' or 'zsh'.
+ `--completion`              |           | Output stern command-line completion code for the specified shell. Can be 'bash', 'zsh' or 'fish'.
  `--container`, `-c`         | `.*`      | Container name when multiple containers in pod. (regular expression)
  `--container-state`         | `running` | Tail containers with state in running, waiting or terminated. To specify multiple states, repeat this or set comma-separated value.
  `--context`                 |           | Kubernetes context to use. Default to current context configured in kubeconfig.
@@ -203,8 +203,8 @@ stern -p
 
 ## Completion
 
-Stern supports command-line auto completion for bash or zsh. `stern
---completion=(bash|zsh)` outputs the shell completion code which work by being
+Stern supports command-line auto completion for bash, zsh or fish. `stern
+--completion=(bash|zsh|fish)` outputs the shell completion code which work by being
 evaluated in `.bashrc`, etc for the specified shell. In addition, Stern
 supports dynamic completion for `--namespace` and `--context`. In order to use
 that, kubectl must be installed on your environment.
@@ -235,6 +235,15 @@ If you use zsh, just source the stern zsh completion code in `.zshrc`.
 source <(stern --completion=zsh)
 ```
 
+if you use fish shell, just source the stern fish completion code.
+
+```sh
+stern --completion=fish | source
+
+# To load completions for each session, execute once:
+stern --completion=fish >~/.config/fish/completions/stern.fish
+```
+
 ## Running with container
 
 You can also use stern using a container:
diff --git a/cmd/cmd.go b/cmd/cmd.go
index 69def35..7c0f240 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -290,7 +290,7 @@ func (o *options) sternConfig() (*stern.Config, error) {
 func (o *options) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVarP(&o.allNamespaces, "all-namespaces", "A", o.allNamespaces, "If present, tail across all namespaces. A specific namespace is ignored even if specified with --namespace.")
 	fs.StringVar(&o.color, "color", o.color, "Force set color output. 'auto':  colorize if tty attached, 'always': always colorize, 'never': never colorize.")
-	fs.StringVar(&o.completion, "completion", o.completion, "Output stern command-line completion code for the specified shell. Can be 'bash' or 'zsh'.")
+	fs.StringVar(&o.completion, "completion", o.completion, "Output stern command-line completion code for the specified shell. Can be 'bash', 'zsh' or 'fish'.")
 	fs.StringVarP(&o.container, "container", "c", o.container, "Container name when multiple containers in pod. (regular expression)")
 	fs.StringSliceVar(&o.containerStates, "container-state", o.containerStates, "Tail containers with state in running, waiting or terminated. To specify multiple states, repeat this or set comma-separated value.")
 	fs.StringVar(&o.context, "context", o.context, "Kubernetes context to use. Default to current context configured in kubeconfig.")
@@ -316,7 +316,7 @@ func (o *options) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVarP(&o.version, "version", "v", o.version, "Print the version and exit.")
 }
 
-func NewSternCmd(stream genericclioptions.IOStreams) *cobra.Command {
+func NewSternCmd(stream genericclioptions.IOStreams) (*cobra.Command, error) {
 	o := NewOptions(stream)
 
 	cmd := &cobra.Command{
@@ -350,9 +350,11 @@ func NewSternCmd(stream genericclioptions.IOStreams) *cobra.Command {
 
 	o.AddFlags(cmd.Flags())
 
-	registerCompletionFunction(cmd)
+	if err := registerCompletionFuncForFlags(cmd, o); err != nil {
+		return cmd, err
+	}
 
-	return cmd
+	return cmd, nil
 }
 
 // makeUnique makes items in string slice unique
diff --git a/cmd/flag_completion.go b/cmd/flag_completion.go
index d038e8d..f88d99c 100644
--- a/cmd/flag_completion.go
+++ b/cmd/flag_completion.go
@@ -16,93 +16,26 @@ package cmd
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
+	"strings"
 
 	"github.com/spf13/cobra"
+	"github.com/stern/stern/kubernetes"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-const (
-	// bash_completion_func depends on kubectl.
-	// If kubectl isn't installed, it doesn't work, but error doesn't occur.
-	bash_completion_func = `
-__is_kubectl_installed=true
-if ! which kubectl >/dev/null 2>&1; then
-	__is_kubectl_installed=false
-fi
-
-__stern_kubectl_override_flag_list=(kubeconfig context namespace)
-__stern_kubectl_override_flags()
-{
-    local ${__stern_kubectl_override_flag_list[*]} two_word_of of
-    for w in "${words[@]}"; do
-        if [ -n "${two_word_of}" ]; then
-            eval "${two_word_of}=\"--${two_word_of}=\${w}\""
-            two_word_of=
-            continue
-        fi
-        for of in "${__stern_kubectl_override_flag_list[@]}"; do
-            case "${w}" in
-                --${of}=*)
-                    eval "${of}=\"${w}\""
-                    ;;
-                --${of})
-                    two_word_of="${of}"
-                    ;;
-            esac
-        done
-        if [ "${w}" == "--all-namespaces" ]; then
-            namespace="--all-namespaces"
-        fi
-    done
-    for of in "${__stern_kubectl_override_flag_list[@]}"; do
-        if eval "test -n \"\$${of}\""; then
-            eval "echo \${${of}}"
-        fi
-    done
-}
-
-__stern_kubectl_get_namespaces()
-{
-    local template kubectl_out
-
-    if ! $__is_kubectl_installed; then
-        return 1
-    fi
-    template="{{ range .items  }}{{ .metadata.name }} {{ end }}"
-    if kubectl_out=$(kubectl get -o template --template="${template}" namespace 2>/dev/null); then
-        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
-    fi
-}
-
-__stern_kubectl_config_get_contexts()
-{
-    local template kubectl_out
-
-    if ! $__is_kubectl_installed; then
-        return 1
-    fi
-    template="{{ range .contexts  }}{{ .name }} {{ end }}"
-    if kubectl_out=$(kubectl config $(__stern_kubectl_override_flags) -o template --template="${template}" view 2>/dev/null); then
-        COMPREPLY=( $( compgen -W "${kubectl_out[*]}" -- "$cur" ) )
-    fi
-}
-`
-)
-
-var bash_completion_flags = map[string]string{
-	"namespace": "__stern_kubectl_get_namespaces",
-	"context":   "__stern_kubectl_config_get_contexts",
-}
-
 func runCompletion(shell string, cmd *cobra.Command, out io.Writer) error {
 	var err error
 
 	switch shell {
 	case "bash":
-		err = runCompletionBash(cmd, out)
+		err = cmd.GenBashCompletion(out)
 	case "zsh":
 		err = runCompletionZsh(cmd, out)
+	case "fish":
+		err = cmd.GenFishCompletion(out, true)
 	default:
 		err = fmt.Errorf("Unsupported shell type: %q", shell)
 	}
@@ -110,171 +43,83 @@ func runCompletion(shell string, cmd *cobra.Command, out io.Writer) error {
 	return err
 }
 
-func runCompletionBash(cmd *cobra.Command, out io.Writer) error {
-	return cmd.GenBashCompletion(out)
-}
-
 // runCompletionZsh is based on `kubectl completion zsh`. This function should
 // be replaced by cobra implementation when cobra itself supports zsh completion.
 // https://github.com/kubernetes/kubernetes/blob/v1.6.1/pkg/kubectl/cmd/completion.go#L136
 func runCompletionZsh(cmd *cobra.Command, out io.Writer) error {
 	b := new(bytes.Buffer)
+	if err := cmd.GenZshCompletion(b); err != nil {
+		return err
+	}
 
-	zshInitialization := `
-__stern_bash_source() {
-    alias shopt=':'
-    alias _expand=_bash_expand
-    alias _complete=_bash_comp
-    emulate -L sh
-    setopt kshglob noshglob braceexpand
-    source "$@"
-}
-__stern_type() {
-    # -t is not supported by zsh
-    if [ "$1" == "-t" ]; then
-        shift
-        # fake Bash 4 to disable "complete -o nospace". Instead
-        # "compopt +-o nospace" is used in the code to toggle trailing
-        # spaces. We don't support that, but leave trailing spaces on
-        # all the time
-        if [ "$1" = "__stern_compopt" ]; then
-            echo builtin
-            return 0
-        fi
-    fi
-    type "$@"
-}
-__stern_compgen() {
-    local completions w
-    completions=( $(compgen "$@") ) || return $?
-    # filter by given word as prefix
-    while [[ "$1" = -* && "$1" != -- ]]; do
-        shift
-        shift
-    done
-    if [[ "$1" == -- ]]; then
-        shift
-    fi
-    for w in "${completions[@]}"; do
-        if [[ "${w}" = "$1"* ]]; then
-            echo "${w}"
-        fi
-    done
-}
-__stern_compopt() {
-    true # don't do anything. Not supported by bashcompinit in zsh
-}
-__stern_ltrim_colon_completions()
-{
-    if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
-        # Remove colon-word prefix from COMPREPLY items
-        local colon_word=${1%${1##*:}}
-        local i=${#COMPREPLY[*]}
-        while [[ $((--i)) -ge 0 ]]; do
-            COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
-        done
-    fi
-}
-__stern_get_comp_words_by_ref() {
-    cur="${COMP_WORDS[COMP_CWORD]}"
-    prev="${COMP_WORDS[${COMP_CWORD}-1]}"
-    words=("${COMP_WORDS[@]}")
-    cword=("${COMP_CWORD[@]}")
-}
-__stern_filedir() {
-    local RET OLD_IFS w qw
-    __debug "_filedir $@ cur=$cur"
-    if [[ "$1" = \~* ]]; then
-        # somehow does not work. Maybe, zsh does not call this at all
-        eval echo "$1"
-        return 0
-    fi
-    OLD_IFS="$IFS"
-    IFS=$'\n'
-    if [ "$1" = "-d" ]; then
-        shift
-        RET=( $(compgen -d) )
-    else
-        RET=( $(compgen -f) )
-    fi
-    IFS="$OLD_IFS"
-    IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"
-    for w in ${RET[@]}; do
-        if [[ ! "${w}" = "${cur}"* ]]; then
-            continue
-        fi
-        if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
-            qw="$(__stern_quote "${w}")"
-            if [ -d "${w}" ]; then
-                COMPREPLY+=("${qw}/")
-            else
-                COMPREPLY+=("${qw}")
-            fi
-        fi
-    done
-}
-__stern_quote() {
-    if [[ $1 == \'* || $1 == \"* ]]; then
-        # Leave out first character
-        printf %q "${1:1}"
-    else
-        printf %q "$1"
-    fi
+	// Cobra doesn't source zsh completion file, explicitly doing it here
+	fmt.Fprintf(b, "compdef _stern stern")
+
+	fmt.Fprint(out, b.String())
+
+	return nil
 }
-autoload -U +X bashcompinit && bashcompinit
-# use word boundary patterns for BSD or GNU sed
-LWORD='[[:<:]]'
-RWORD='[[:>:]]'
-if sed --help 2>&1 | grep -q GNU; then
-    LWORD='\<'
-    RWORD='\>'
-fi
-__stern_convert_bash_to_zsh() {
-    sed \
-    -e 's/declare -F/whence -w/' \
-    -e 's/_get_comp_words_by_ref "\$@"/_get_comp_words_by_ref "\$*"/' \
-    -e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
-    -e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
-    -e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
-    -e "s/${LWORD}_filedir${RWORD}/__stern_filedir/g" \
-    -e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__stern_get_comp_words_by_ref/g" \
-    -e "s/${LWORD}__ltrim_colon_completions${RWORD}/__stern_ltrim_colon_completions/g" \
-    -e "s/${LWORD}compgen${RWORD}/__stern_compgen/g" \
-    -e "s/${LWORD}compopt${RWORD}/__stern_compopt/g" \
-    -e "s/${LWORD}declare${RWORD}/builtin declare/g" \
-    -e "s/\\\$(type${RWORD}/\$(__stern_type/g" \
-    <<'BASH_COMPLETION_EOF'
-`
-	b.Write([]byte(zshInitialization))
 
-	if err := cmd.GenBashCompletion(b); err != nil {
+func registerCompletionFuncForFlags(cmd *cobra.Command, o *options) error {
+	if err := cmd.RegisterFlagCompletionFunc("namespace", namespaceCompletionFunc(o)); err != nil {
 		return err
 	}
 
-	zshTail := `
-BASH_COMPLETION_EOF
+	if err := cmd.RegisterFlagCompletionFunc("context", contextCompletionFunc(o)); err != nil {
+		return err
+	}
+
+	return nil
 }
-__stern_bash_source <(__stern_convert_bash_to_zsh)
-`
-	b.Write([]byte(zshTail))
 
-	fmt.Fprintln(out, b.String())
+// namespaceCompletionFunc is a completion function that completes namespaces
+// that match the toComplete prefix.
+func namespaceCompletionFunc(o *options) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		clientConfig := kubernetes.NewClientConfig(o.kubeConfig, o.context)
+		clientset, err := kubernetes.NewClientSet(clientConfig)
+		if err != nil {
+			return compError(err)
+		}
 
-	return nil
+		namespaceList, err := clientset.CoreV1().Namespaces().List(context.TODO(), metav1.ListOptions{})
+		if err != nil {
+			return compError(err)
+		}
+
+		var comps []string
+		for _, ns := range namespaceList.Items {
+			if strings.HasPrefix(ns.GetName(), toComplete) {
+				comps = append(comps, ns.GetName())
+			}
+		}
+
+		return comps, cobra.ShellCompDirectiveNoFileComp
+	}
 }
 
-func registerCompletionFunction(cmd *cobra.Command) {
-	// Specify custom bash completion function
-	cmd.BashCompletionFunction = bash_completion_func
-	for name, completion := range bash_completion_flags {
-		if cmd.Flag(name) != nil {
-			if cmd.Flag(name).Annotations == nil {
-				cmd.Flag(name).Annotations = map[string][]string{}
+// contextCompletionFunc is a completion function that completes contexts
+// that match the toComplete prefix.
+func contextCompletionFunc(o *options) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		clientConfig := kubernetes.NewClientConfig(o.kubeConfig, o.context)
+		config, err := clientConfig.RawConfig()
+		if err != nil {
+			return compError(err)
+		}
+
+		var comps []string
+		for name := range config.Contexts {
+			if strings.HasPrefix(name, toComplete) {
+				comps = append(comps, name)
 			}
-			cmd.Flag(name).Annotations[cobra.BashCompCustom] = append(
-				cmd.Flag(name).Annotations[cobra.BashCompCustom],
-				completion,
-			)
 		}
+
+		return comps, cobra.ShellCompDirectiveNoFileComp
 	}
 }
+
+func compError(err error) ([]string, cobra.ShellCompDirective) {
+	cobra.CompError(err.Error())
+	return nil, cobra.ShellCompDirectiveError
+}
diff --git a/main.go b/main.go
index 05c6fc2..a8dfab3 100644
--- a/main.go
+++ b/main.go
@@ -15,6 +15,7 @@
 package main
 
 import (
+	"log"
 	"os"
 
 	"github.com/stern/stern/cmd"
@@ -23,7 +24,12 @@ import (
 
 func main() {
 	streams := genericclioptions.IOStreams{Out: os.Stdout, ErrOut: os.Stderr}
-	if err := cmd.NewSternCmd(streams).Execute(); err != nil {
+	stern, err := cmd.NewSternCmd(streams)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	if err := stern.Execute(); err != nil {
 		os.Exit(1)
 	}
 }


Test patch:
diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go
index ffea495..e6d6ce5 100644
--- a/cmd/cmd_test.go
+++ b/cmd/cmd_test.go
@@ -28,10 +28,13 @@ func TestSternCommand(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			streams, _, out, _ := genericclioptions.NewTestIOStreams()
-			cmd := NewSternCmd(streams)
-			cmd.SetArgs(tt.args)
+			stern, err := NewSternCmd(streams)
+			if err != nil {
+				t.Fatal(err)
+			}
+			stern.SetArgs(tt.args)
 
-			if err := cmd.Execute(); err != nil {
+			if err := stern.Execute(); err != nil {
 				t.Fatal(err)
 			}