higress/hgctl/pkg/completion.go

// Copyright (c) 2022 Alibaba Group Holding Ltd.
//
// 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,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package hgctl

import (
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/spf13/cobra"
)

const completionDesc = `
Generate autocompletion scripts for hgctl for the specified shell.
`

const bashCompDesc = `
Generate the autocompletion script for the bash shell.

This script depends on the 'bash-completion' package.
If it is not installed already, you can install it via your OS's package manager.

To load completions in your current shell session:

        source <(hgctl completion bash)

To load completions for every new session, execute once:

#### Linux:

        hgctl completion bash > /etc/bash_completion.d/hgctl

#### macOS:

        hgctl completion bash > $(brew --prefix)/etc/bash_completion.d/hgctl

You will need to start a new shell for this setup to take effect.
`

const zshCompDesc = `
Generate the autocompletion script for the zsh shell.

If shell completion is not already enabled in your environment you will need
to enable it.  You can execute the following once:

        echo "autoload -U compinit; compinit" >> ~/.zshrc

To load completions in your current shell session:

        source <(hgctl completion zsh); compdef _hgctl hgctl

To load completions for every new session, execute once:

#### Linux:

        hgctl completion zsh > "${fpath[1]}/_hgctl"

#### macOS:

        hgctl completion zsh > $(brew --prefix)/share/zsh/site-functions/_hgctl

You will need to start a new shell for this setup to take effect.
`

const fishCompDesc = `
Generate the autocompletion script for the fish shell.

To load completions in your current shell session:

        hgctl completion fish | source

To load completions for every new session, execute once:

        hgctl completion fish > ~/.config/fish/completions/hgctl.fish

You will need to start a new shell for this setup to take effect.
`

const powershellCompDesc = `
Generate the autocompletion script for powershell.

To load completions in your current shell session:

        hgctl completion powershell | Out-String | Invoke-Expression

To load completions for every new session, add the output of the above command
to your powershell profile.
`

const (
	noDescFlagName = "no-descriptions"
	noDescFlagText = "disable completion descriptions"
)

var disableCompDescriptions bool

// newCompletionCmd creates a new completion command for hgctl
func newCompletionCmd(out io.Writer) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "completion",
		Short: "generate autocompletion scripts for the specified shell",
		Long:  completionDesc,
		Args:  cobra.NoArgs,
	}

	bash := &cobra.Command{
		Use:               "bash",
		Short:             "generate autocompletion script for bash",
		Long:              bashCompDesc,
		Args:              cobra.NoArgs,
		ValidArgsFunction: noCompletions,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletionBash(out, cmd)
		},
	}
	bash.Flags().BoolVar(&disableCompDescriptions, noDescFlagName, false, noDescFlagText)

	zsh := &cobra.Command{
		Use:               "zsh",
		Short:             "generate autocompletion script for zsh",
		Long:              zshCompDesc,
		Args:              cobra.NoArgs,
		ValidArgsFunction: noCompletions,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletionZsh(out, cmd)
		},
	}
	zsh.Flags().BoolVar(&disableCompDescriptions, noDescFlagName, false, noDescFlagText)

	fish := &cobra.Command{
		Use:               "fish",
		Short:             "generate autocompletion script for fish",
		Long:              fishCompDesc,
		Args:              cobra.NoArgs,
		ValidArgsFunction: noCompletions,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletionFish(out, cmd)
		},
	}
	fish.Flags().BoolVar(&disableCompDescriptions, noDescFlagName, false, noDescFlagText)

	powershell := &cobra.Command{
		Use:               "powershell",
		Short:             "generate autocompletion script for powershell",
		Long:              powershellCompDesc,
		Args:              cobra.NoArgs,
		ValidArgsFunction: noCompletions,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runCompletionPowershell(out, cmd)
		},
	}
	powershell.Flags().BoolVar(&disableCompDescriptions, noDescFlagName, false, noDescFlagText)

	cmd.AddCommand(bash, zsh, fish, powershell)

	return cmd
}

func runCompletionBash(out io.Writer, cmd *cobra.Command) error {
	err := cmd.Root().GenBashCompletionV2(out, !disableCompDescriptions)

	// In case the user renamed the hgctl binary, we hook the new binary name to the completion function
	if binary := filepath.Base(os.Args[0]); binary != "hgctl" {
		renamedBinaryHook := `
# Hook the command used to generate the completion script
# to the hgctl completion function to handle the case where
# the user renamed the hgctl binary
if [[ $(type -t compopt) = "builtin" ]]; then
    complete -o default -F __start_hgctl %[1]s
else
    complete -o default -o nospace -F __start_hgctl %[1]s
fi
`
		fmt.Fprintf(out, renamedBinaryHook, binary)
	}

	return err
}

func runCompletionZsh(out io.Writer, cmd *cobra.Command) error {
	var err error
	if disableCompDescriptions {
		err = cmd.Root().GenZshCompletionNoDesc(out)
	} else {
		err = cmd.Root().GenZshCompletion(out)
	}

	// In case the user renamed the hgctl binary, we hook the new binary name to the completion function
	if binary := filepath.Base(os.Args[0]); binary != "hgctl" {
		renamedBinaryHook := `
# Hook the command used to generate the completion script
# to the hgctl completion function to handle the case where
# the user renamed the hgctl binary
compdef _hgctl %[1]s
`
		fmt.Fprintf(out, renamedBinaryHook, binary)
	}

	// Cobra doesn't source zsh completion file, explicitly doing it here
	fmt.Fprintf(out, "compdef _hgctl hgctl")

	return err
}

func runCompletionFish(out io.Writer, cmd *cobra.Command) error {
	return cmd.Root().GenFishCompletion(out, !disableCompDescriptions)
}

func runCompletionPowershell(out io.Writer, cmd *cobra.Command) error {
	if disableCompDescriptions {
		return cmd.Root().GenPowerShellCompletion(out)
	}
	return cmd.Root().GenPowerShellCompletionWithDesc(out)
}

// Function to disable file completion
func noCompletions(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
	return nil, cobra.ShellCompDirectiveNoFileComp
}