You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

130 lines
3.6 KiB

8 years ago
8 years ago
8 years ago
8 years ago
  1. package cli
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "os"
  7. "path/filepath"
  8. "github.com/spf13/cobra"
  9. )
  10. // WriteConfigVals writes a toml file with the given values.
  11. // It returns an error if writing was impossible.
  12. func WriteConfigVals(dir string, vals map[string]string) error {
  13. data := ""
  14. for k, v := range vals {
  15. data += fmt.Sprintf("%s = \"%s\"\n", k, v)
  16. }
  17. cfile := filepath.Join(dir, "config.toml")
  18. return os.WriteFile(cfile, []byte(data), 0600)
  19. }
  20. // RunWithArgs executes the given command with the specified command line args
  21. // and environmental variables set. It returns any error returned from cmd.Execute()
  22. func RunWithArgs(cmd Executable, args []string, env map[string]string) error {
  23. oargs := os.Args
  24. oenv := map[string]string{}
  25. // defer returns the environment back to normal
  26. defer func() {
  27. os.Args = oargs
  28. for k, v := range oenv {
  29. os.Setenv(k, v)
  30. }
  31. }()
  32. // set the args and env how we want them
  33. os.Args = args
  34. for k, v := range env {
  35. // backup old value if there, to restore at end
  36. oenv[k] = os.Getenv(k)
  37. err := os.Setenv(k, v)
  38. if err != nil {
  39. return err
  40. }
  41. }
  42. // and finally run the command
  43. return cmd.Execute()
  44. }
  45. // RunCaptureWithArgs executes the given command with the specified command
  46. // line args and environmental variables set. It returns string fields
  47. // representing output written to stdout and stderr, additionally any error
  48. // from cmd.Execute() is also returned
  49. func RunCaptureWithArgs(cmd Executable, args []string, env map[string]string) (stdout, stderr string, err error) {
  50. oldout, olderr := os.Stdout, os.Stderr // keep backup of the real stdout
  51. rOut, wOut, _ := os.Pipe()
  52. rErr, wErr, _ := os.Pipe()
  53. os.Stdout, os.Stderr = wOut, wErr
  54. defer func() {
  55. os.Stdout, os.Stderr = oldout, olderr // restoring the real stdout
  56. }()
  57. // copy the output in a separate goroutine so printing can't block indefinitely
  58. copyStd := func(reader *os.File) *(chan string) {
  59. stdC := make(chan string)
  60. go func() {
  61. var buf bytes.Buffer
  62. // io.Copy will end when we call reader.Close() below
  63. io.Copy(&buf, reader) //nolint:errcheck //ignore error
  64. select {
  65. case <-cmd.Context().Done():
  66. case stdC <- buf.String():
  67. }
  68. }()
  69. return &stdC
  70. }
  71. outC := copyStd(rOut)
  72. errC := copyStd(rErr)
  73. // now run the command
  74. err = RunWithArgs(cmd, args, env)
  75. // and grab the stdout to return
  76. wOut.Close()
  77. wErr.Close()
  78. stdout = <-*outC
  79. stderr = <-*errC
  80. return stdout, stderr, err
  81. }
  82. // NewCompletionCmd returns a cobra.Command that generates bash and zsh
  83. // completion scripts for the given root command. If hidden is true, the
  84. // command will not show up in the root command's list of available commands.
  85. func NewCompletionCmd(rootCmd *cobra.Command, hidden bool) *cobra.Command {
  86. flagZsh := "zsh"
  87. cmd := &cobra.Command{
  88. Use: "completion",
  89. Short: "Generate shell completion scripts",
  90. Long: fmt.Sprintf(`Generate Bash and Zsh completion scripts and print them to STDOUT.
  91. Once saved to file, a completion script can be loaded in the shell's
  92. current session as shown:
  93. $ . <(%s completion)
  94. To configure your bash shell to load completions for each session add to
  95. your $HOME/.bashrc or $HOME/.profile the following instruction:
  96. . <(%s completion)
  97. `, rootCmd.Use, rootCmd.Use),
  98. RunE: func(cmd *cobra.Command, _ []string) error {
  99. zsh, err := cmd.Flags().GetBool(flagZsh)
  100. if err != nil {
  101. return err
  102. }
  103. if zsh {
  104. return rootCmd.GenZshCompletion(cmd.OutOrStdout())
  105. }
  106. return rootCmd.GenBashCompletion(cmd.OutOrStdout())
  107. },
  108. Hidden: hidden,
  109. Args: cobra.NoArgs,
  110. }
  111. cmd.Flags().Bool(flagZsh, false, "Generate Zsh completion script")
  112. return cmd
  113. }