powershell_completions.go 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  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. package cobra
  7. import (
  8. "bytes"
  9. "fmt"
  10. "io"
  11. "os"
  12. "strings"
  13. "github.com/spf13/pflag"
  14. )
  15. var powerShellCompletionTemplate = `using namespace System.Management.Automation
  16. using namespace System.Management.Automation.Language
  17. Register-ArgumentCompleter -Native -CommandName '%s' -ScriptBlock {
  18. param($wordToComplete, $commandAst, $cursorPosition)
  19. $commandElements = $commandAst.CommandElements
  20. $command = @(
  21. '%s'
  22. for ($i = 1; $i -lt $commandElements.Count; $i++) {
  23. $element = $commandElements[$i]
  24. if ($element -isnot [StringConstantExpressionAst] -or
  25. $element.StringConstantType -ne [StringConstantType]::BareWord -or
  26. $element.Value.StartsWith('-')) {
  27. break
  28. }
  29. $element.Value
  30. }
  31. ) -join ';'
  32. $completions = @(switch ($command) {%s
  33. })
  34. $completions.Where{ $_.CompletionText -like "$wordToComplete*" } |
  35. Sort-Object -Property ListItemText
  36. }`
  37. func generatePowerShellSubcommandCases(out io.Writer, cmd *Command, previousCommandName string) {
  38. var cmdName string
  39. if previousCommandName == "" {
  40. cmdName = cmd.Name()
  41. } else {
  42. cmdName = fmt.Sprintf("%s;%s", previousCommandName, cmd.Name())
  43. }
  44. fmt.Fprintf(out, "\n '%s' {", cmdName)
  45. cmd.Flags().VisitAll(func(flag *pflag.Flag) {
  46. if nonCompletableFlag(flag) {
  47. return
  48. }
  49. usage := escapeStringForPowerShell(flag.Usage)
  50. if len(flag.Shorthand) > 0 {
  51. fmt.Fprintf(out, "\n [CompletionResult]::new('-%s', '%s', [CompletionResultType]::ParameterName, '%s')", flag.Shorthand, flag.Shorthand, usage)
  52. }
  53. fmt.Fprintf(out, "\n [CompletionResult]::new('--%s', '%s', [CompletionResultType]::ParameterName, '%s')", flag.Name, flag.Name, usage)
  54. })
  55. for _, subCmd := range cmd.Commands() {
  56. usage := escapeStringForPowerShell(subCmd.Short)
  57. fmt.Fprintf(out, "\n [CompletionResult]::new('%s', '%s', [CompletionResultType]::ParameterValue, '%s')", subCmd.Name(), subCmd.Name(), usage)
  58. }
  59. fmt.Fprint(out, "\n break\n }")
  60. for _, subCmd := range cmd.Commands() {
  61. generatePowerShellSubcommandCases(out, subCmd, cmdName)
  62. }
  63. }
  64. func escapeStringForPowerShell(s string) string {
  65. return strings.Replace(s, "'", "''", -1)
  66. }
  67. // GenPowerShellCompletion generates PowerShell completion file and writes to the passed writer.
  68. func (c *Command) GenPowerShellCompletion(w io.Writer) error {
  69. buf := new(bytes.Buffer)
  70. var subCommandCases bytes.Buffer
  71. generatePowerShellSubcommandCases(&subCommandCases, c, "")
  72. fmt.Fprintf(buf, powerShellCompletionTemplate, c.Name(), c.Name(), subCommandCases.String())
  73. _, err := buf.WriteTo(w)
  74. return err
  75. }
  76. // GenPowerShellCompletionFile generates PowerShell completion file.
  77. func (c *Command) GenPowerShellCompletionFile(filename string) error {
  78. outFile, err := os.Create(filename)
  79. if err != nil {
  80. return err
  81. }
  82. defer outFile.Close()
  83. return c.GenPowerShellCompletion(outFile)
  84. }