cmd.go 1.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. package app
  2. import (
  3. "dy-admin/pkg/comopts"
  4. "fmt"
  5. "github.com/fatih/color"
  6. "github.com/spf13/cobra"
  7. "os"
  8. )
  9. // Command 应用子命令结构
  10. type Command struct {
  11. usage string
  12. desc string
  13. options comopts.CommandLineOptions
  14. commands []*Command
  15. runFunc RunCommandFunc
  16. }
  17. type RunCommandFunc func(args []string) error
  18. type CommandOption func(*Command)
  19. func WithCommandOptions(opt comopts.CommandLineOptions) CommandOption {
  20. return func(command *Command) {
  21. command.options = opt
  22. }
  23. }
  24. func WithCommandRunFunc(run RunCommandFunc) CommandOption {
  25. return func(c *Command) {
  26. c.runFunc = run
  27. }
  28. }
  29. func NewCommand(usage, desc string, opts ...CommandOption) *Command {
  30. c := &Command{
  31. usage: usage,
  32. desc: desc,
  33. }
  34. for _, o := range opts {
  35. o(c)
  36. }
  37. return c
  38. }
  39. // AddCommand 当前命令添加字命令
  40. func (c *Command) AddCommand(cmd *Command) {
  41. c.commands = append(c.commands, cmd)
  42. }
  43. // AddCommands 当前命令添加字命令
  44. func (c *Command) AddCommands(cmds ...*Command) {
  45. c.commands = append(c.commands, cmds...)
  46. }
  47. func (c *Command) cobraCommand() *cobra.Command {
  48. cmd := &cobra.Command{
  49. Use: c.usage,
  50. Short: c.desc,
  51. }
  52. cmd.SetOut(os.Stdout)
  53. cmd.Flags().SortFlags = false
  54. if len(c.commands) > 0 {
  55. for _, command := range c.commands {
  56. cmd.AddCommand(command.cobraCommand())
  57. }
  58. }
  59. if c.runFunc != nil {
  60. cmd.Run = c.runCommand
  61. }
  62. if c.options != nil {
  63. for _, f := range c.options.Flags().FlagSets {
  64. cmd.Flags().AddFlagSet(f)
  65. }
  66. // c.comopts.AddFlags(cmd.Flags())
  67. }
  68. addHelpCommandFlag(c.usage, cmd.Flags())
  69. return cmd
  70. }
  71. func (c *Command) runCommand(cmd *cobra.Command, args []string) {
  72. if c.runFunc != nil {
  73. if err := c.runFunc(args); err != nil {
  74. fmt.Printf("%v %v\n", color.RedString("Error:"), err)
  75. os.Exit(1)
  76. }
  77. }
  78. }