command.go

Documentation: edge-infra.dev/pkg/lib/cli/command

     1  package command
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"strings"
     8  	"text/tabwriter"
     9  
    10  	"github.com/peterbourgon/ff/v3"
    11  	"github.com/peterbourgon/ff/v3/ffcli"
    12  
    13  	"edge-infra.dev/pkg/lib/cli/rags"
    14  )
    15  
    16  var (
    17  	options = []ff.Option{ff.WithEnvVarNoPrefix()}
    18  )
    19  
    20  // Command is a simple ffcli.Command wrapper used to make consistent ffcli Commands
    21  // with helpers.
    22  type Command struct {
    23  	// Exec is the function executed when the command is called. It is wrapped
    24  	// in Command.WrappedExec so that the command's context can be decorated with
    25  	// CLI state based on the command being executed.
    26  	Exec func(ctx context.Context, args []string) error
    27  
    28  	// CLI machinery configuration.
    29  	ShortUsage string // One-line usage message. Command name is parsed from this
    30  	ShortHelp  string
    31  	LongHelp   string
    32  	Rags       *rags.RagSet
    33  	Flags      []*rags.Rag                 // Flags specific to this command.
    34  	Options    []ff.Option                 // Options specific to this command
    35  	Commands   []*Command                  // Subcommands
    36  	Usage      func(*ffcli.Command) string // Usage dictates how the cli is displayed
    37  	Extensions []Extension
    38  
    39  	help bool // --help
    40  }
    41  
    42  // TODO: provide ability to store extension in context similar to test/f2/fctx?
    43  
    44  type Extension interface {
    45  	RegisterFlags(rs *rags.RagSet)
    46  }
    47  
    48  type AfterParser interface {
    49  	AfterParse() error
    50  }
    51  
    52  func (c *Command) LongName() string {
    53  	n := c.ShortUsage
    54  	if i := strings.Index(n, " ["); i >= 0 {
    55  		n = n[:i]
    56  	}
    57  	if i := strings.Index(n, " "); i >= 0 {
    58  		return n[i+1:]
    59  	}
    60  	return ""
    61  }
    62  
    63  func (c *Command) Name() string {
    64  	n := c.LongName()
    65  	if i := strings.LastIndex(n, " "); i >= 0 {
    66  		n = n[i+1:]
    67  	}
    68  	return n
    69  }
    70  
    71  func (c *Command) Command() *ffcli.Command {
    72  	c.Rags = rags.New(c.ShortUsage, flag.ContinueOnError)
    73  	c.registerFlags()
    74  
    75  	if c.Exec == nil {
    76  		c.Exec = func(context.Context, []string) error {
    77  			return nil
    78  		}
    79  		// root Exec will print the usage info same as --help would
    80  		c.help = true
    81  	}
    82  
    83  	cmd := &ffcli.Command{
    84  		Name:       c.Name(),
    85  		ShortUsage: c.ShortUsage,
    86  		ShortHelp:  c.ShortHelp,
    87  		LongHelp:   c.LongHelp,
    88  		FlagSet:    c.Rags.FlagSet(),
    89  		Options:    append(options, c.Options...),
    90  		Exec:       c.wExec(),
    91  		UsageFunc:  c.Usage,
    92  	}
    93  	for _, s := range c.Commands {
    94  		cmd.Subcommands = append(cmd.Subcommands, s.Command())
    95  	}
    96  	cmd.UsageFunc = func(_ *ffcli.Command) string {
    97  		return c.DefaultUsage()
    98  	}
    99  	return cmd
   100  }
   101  
   102  func (c *Command) afterParse(ctx context.Context) (context.Context, error) {
   103  	for _, e := range c.Extensions {
   104  		if ap, ok := e.(AfterParser); ok {
   105  			if err := ap.AfterParse(); err != nil {
   106  				return ctx, err
   107  			}
   108  		}
   109  	}
   110  
   111  	return ctx, nil
   112  }
   113  
   114  // wExec wraps the configured Exec function to implement common command logic
   115  // and capture command state.
   116  func (c *Command) wExec() func(context.Context, []string) error {
   117  	return (func() func(context.Context, []string) error {
   118  		cmd := c
   119  		return func(ctx context.Context, args []string) error {
   120  			if cmd.help {
   121  				fmt.Println(c.DefaultUsage())
   122  				return nil
   123  			}
   124  
   125  			ctx, err := cmd.afterParse(ctx)
   126  			if err != nil {
   127  				return err
   128  			}
   129  
   130  			return cmd.Exec(ctx, args)
   131  		}
   132  	})()
   133  }
   134  
   135  func (c *Command) registerFlags() {
   136  	// TODO: use declarative rags flag for this
   137  	c.Rags.BoolVar(&c.help, "help", false, "display help information")
   138  	for _, e := range c.Extensions {
   139  		e.RegisterFlags(c.Rags)
   140  	}
   141  	c.Rags.Add(c.Flags...)
   142  }
   143  
   144  func (c *Command) DefaultUsage() string {
   145  	var b strings.Builder
   146  	fmt.Fprintf(&b, "Usage:\n")
   147  	if c.ShortUsage != "" {
   148  		fmt.Fprintf(&b, "  %s\n", c.ShortUsage)
   149  	} else {
   150  		fmt.Fprintf(&b, "  %s\n", c.Name())
   151  	}
   152  	fmt.Fprintf(&b, "\n")
   153  
   154  	if c.LongHelp != "" {
   155  		fmt.Fprintf(&b, "%s\n\n", c.LongHelp)
   156  	}
   157  
   158  	if len(c.Commands) > 0 {
   159  		fmt.Fprintf(&b, "Subcommands:\n")
   160  		tw := tabwriter.NewWriter(&b, 0, 2, 2, ' ', 0)
   161  		for _, subcommand := range c.Commands {
   162  			fmt.Fprintf(tw, "  %s\t%s\n", subcommand.Name(), subcommand.ShortHelp)
   163  		}
   164  		tw.Flush()
   165  		fmt.Fprintf(&b, "\n")
   166  	}
   167  	c.Rags.SetOutput(&b)
   168  	c.Rags.Usage()
   169  	return (strings.TrimSpace(b.String()) + "\n")
   170  }
   171
View as plain text