1 // PowerShell completions are based on the amazing work from clap:

     2 // https://github.com/clap-rs/clap/blob/3294d18efe5f264d12c9035f404c7d189d4824e1/src/completions/powershell.rs

     3 //

     4 // The generated scripts require PowerShell v5.0+ (which comes Windows 10, but

     5 // can be downloaded separately for windows 7 or 8.1).

     6 

     7 package cobra

     8 

     9 import (

    10 	"bytes"

    11 	"fmt"

    12 	"io"

    13 	"os"

    14 	"strings"

    15 

    16 	"github.com/spf13/pflag"

    17 )

    18 

    19 var powerShellCompletionTemplate = `using namespace System.Management.Automation

    20 using namespace System.Management.Automation.Language

    21 Register-ArgumentCompleter -Native -CommandName '%s' -ScriptBlock {

    22     param($wordToComplete, $commandAst, $cursorPosition)

    23     $commandElements = $commandAst.CommandElements

    24     $command = @(

    25         '%s'

    26         for ($i = 1; $i -lt $commandElements.Count; $i++) {

    27             $element = $commandElements[$i]

    28             if ($element -isnot [StringConstantExpressionAst] -or

    29                 $element.StringConstantType -ne [StringConstantType]::BareWord -or

    30                 $element.Value.StartsWith('-')) {

    31                 break

    32             }

    33             $element.Value

    34         }

    35     ) -join ';'

    36     $completions = @(switch ($command) {%s

    37     })

    38     $completions.Where{ $_.CompletionText -like "$wordToComplete*" } |

    39         Sort-Object -Property ListItemText

    40 }`

    41 

    42 func generatePowerShellSubcommandCases(out io.Writer, cmd *Command, previousCommandName string) {

    43 	var cmdName string

    44 	if previousCommandName == "" {

    45 		cmdName = cmd.Name()

    46 	} else {

    47 		cmdName = fmt.Sprintf("%s;%s", previousCommandName, cmd.Name())

    48 	}

    49 

    50 	fmt.Fprintf(out, "\n        '%s' {", cmdName)

    51 

    52 	cmd.Flags().VisitAll(func(flag *pflag.Flag) {

    53 		if nonCompletableFlag(flag) {

    54 			return

    55 		}

    56 		usage := escapeStringForPowerShell(flag.Usage)

    57 		if len(flag.Shorthand) > 0 {

    58 			fmt.Fprintf(out, "\n            [CompletionResult]::new('-%s', '%s', [CompletionResultType]::ParameterName, '%s')", flag.Shorthand, flag.Shorthand, usage)

    59 		}

    60 		fmt.Fprintf(out, "\n            [CompletionResult]::new('--%s', '%s', [CompletionResultType]::ParameterName, '%s')", flag.Name, flag.Name, usage)

    61 	})

    62 

    63 	for _, subCmd := range cmd.Commands() {

    64 		usage := escapeStringForPowerShell(subCmd.Short)

    65 		fmt.Fprintf(out, "\n            [CompletionResult]::new('%s', '%s', [CompletionResultType]::ParameterValue, '%s')", subCmd.Name(), subCmd.Name(), usage)

    66 	}

    67 

    68 	fmt.Fprint(out, "\n            break\n        }")

    69 

    70 	for _, subCmd := range cmd.Commands() {

    71 		generatePowerShellSubcommandCases(out, subCmd, cmdName)

    72 	}

    73 }

    74 

    75 func escapeStringForPowerShell(s string) string {

    76 	return strings.Replace(s, "'", "''", -1)

    77 }

    78 

    79 // GenPowerShellCompletion generates PowerShell completion file and writes to the passed writer.

    80 func (c *Command) GenPowerShellCompletion(w io.Writer) error {

    81 	buf := new(bytes.Buffer)

    82 

    83 	var subCommandCases bytes.Buffer

    84 	generatePowerShellSubcommandCases(&subCommandCases, c, "")

    85 	fmt.Fprintf(buf, powerShellCompletionTemplate, c.Name(), c.Name(), subCommandCases.String())

    86 

    87 	_, err := buf.WriteTo(w)

    88 	return err

    89 }

    90 

    91 // GenPowerShellCompletionFile generates PowerShell completion file.

    92 func (c *Command) GenPowerShellCompletionFile(filename string) error {

    93 	outFile, err := os.Create(filename)

    94 	if err != nil {

    95 		return err

    96 	}

    97 	defer outFile.Close()

    98 

    99 	return c.GenPowerShellCompletion(outFile)

   100 }