shorthand flags + cleanup

Author: ethanndickson

command.go

@@ -620,14 +620,7 @@ func (inv *Invocation) with(fn func(*Invocation)) *Invocation {
 }
 
 func (inv *Invocation) flagHandler(word string) bool {
-	opt := inv.Command.Options.ByFlag(word[2:])
-	if opt != nil && opt.CompletionHandler != nil {
-		for _, e := range opt.CompletionHandler(inv) {
-			fmt.Fprintf(inv.Stdout, "%s\n", e)
-		}
-		return true
-	}
-	return false
+	return inv.doFlagCompletion("", word)
 }
 
 func (inv *Invocation) equalsFlagHandler(word string) bool {
@@ -638,11 +631,26 @@ func (inv *Invocation) equalsFlagHandler(word string) bool {
 	} else {
 		inv.CurWord = ""
 	}
-	outPrefix := word + "="
+	prefix := word + "="
+	return inv.doFlagCompletion(prefix, word)
+}
+
+func (inv *Invocation) doFlagCompletion(prefix, word string) bool {
 	opt := inv.Command.Options.ByFlag(word[2:])
-	if opt != nil && opt.CompletionHandler != nil {
-		for _, e := range opt.CompletionHandler(inv) {
-			fmt.Fprintf(inv.Stdout, "%s%s\n", outPrefix, e)
+	if opt == nil {
+		return false
+	}
+	if opt.CompletionHandler != nil {
+		completions := opt.CompletionHandler(inv)
+		for _, completion := range completions {
+			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, completion)
+		}
+		return true
+	}
+	val, ok := opt.Value.(*Enum)
+	if ok {
+		for _, choice := range val.Choices {
+			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, choice)
 		}
 		return true
 	}

command_test.go

@@ -64,14 +64,16 @@ func SampleCommand(t *testing.T) *serpent.Command {
 				Short: "Example with required flags",
 				Options: serpent.OptionSet{
 					serpent.Option{
-						Name:     "req-bool",
-						Flag:     "req-bool",
-						Value:    serpent.BoolOf(&reqBool),
-						Required: true,
+						Name:          "req-bool",
+						Flag:          "req-bool",
+						FlagShorthand: "b",
+						Value:         serpent.BoolOf(&reqBool),
+						Required:      true,
 					},
 					serpent.Option{
-						Name: "req-string",
-						Flag: "req-string",
+						Name:          "req-string",
+						Flag:          "req-string",
+						FlagShorthand: "s",
 						Value: serpent.Validate(serpent.StringOf(&reqStr), func(value *serpent.String) error {
 							ok := strings.Contains(value.String(), " ")
 							if !ok {
@@ -82,15 +84,15 @@ func SampleCommand(t *testing.T) *serpent.Command {
 						Required: true,
 					},
 					serpent.Option{
-						Name:              "req-enum",
-						Flag:              "req-enum",
-						Value:             serpent.EnumOf(&enumStr, enumChoices...),
-						CompletionHandler: completion.EnumHandler(enumChoices...),
+						Name:  "req-enum",
+						Flag:  "req-enum",
+						Value: serpent.EnumOf(&enumStr, enumChoices...),
 					},
 					serpent.Option{
-						Name:  "req-array",
-						Flag:  "req-array",
-						Value: serpent.StringArrayOf(&reqArr),
+						Name:          "req-array",
+						Flag:          "req-array",
+						FlagShorthand: "a",
+						Value:         serpent.StringArrayOf(&reqArr),
 					},
 				},
 				HelpHandler: func(i *serpent.Invocation) error {

completion/all.go

@@ -30,10 +30,6 @@ func ShellOptions(choice *string) *serpent.Enum {
 	return serpent.EnumOf(choice, BashShell, FishShell, ZShell, Powershell)
 }
 
-func ShellHandler() serpent.CompletionHandlerFunc {
-	return EnumHandler(BashShell, FishShell, ZShell, Powershell)
-}
-
 func GetCompletion(writer io.Writer, shell string, cmdName string) error {
 	fn, ok := shellCompletionByName[shell]
 	if !ok {

completion/handlers.go

@@ -9,12 +9,6 @@ import (
 	"github.com/coder/serpent"
 )
 
-func EnumHandler(choices ...string) serpent.CompletionHandlerFunc {
-	return func(inv *serpent.Invocation) []string {
-		return choices
-	}
-}
-
 // FileHandler returns a handler that completes files, using the
 // given filter func, which may be nil.
 func FileHandler(filter func(info os.FileInfo) bool) serpent.CompletionHandlerFunc {

completion_test.go

@@ -25,6 +25,16 @@ func TestCompletion(t *testing.T) {
 		require.Equal(t, "altfile\nfile\nrequired-flag\ntoupper\n--prefix\n--verbose\n", io.Stdout.String())
 	})
 
+	t.Run("SubcommandNoPartial", func(t *testing.T) {
+		t.Parallel()
+		i := cmd().Invoke("f")
+		i.Environ.Set(serpent.CompletionModeEnv, "1")
+		io := fakeIO(i)
+		err := i.Run()
+		require.NoError(t, err)
+		require.Equal(t, "altfile\nfile\nrequired-flag\ntoupper\n--prefix\n--verbose\n", io.Stdout.String())
+	})
+
 	t.Run("SubcommandComplete", func(t *testing.T) {
 		t.Parallel()
 		i := cmd().Invoke("required-flag")
@@ -47,7 +57,17 @@ func TestCompletion(t *testing.T) {
 
 	t.Run("FlagExhaustive", func(t *testing.T) {
 		t.Parallel()
-		i := cmd().Invoke("required-flag", "--req-bool", "--req-string", "foo bar", "--req-array", "asdf")
+		i := cmd().Invoke("required-flag", "--req-bool", "--req-string", "foo bar", "--req-array", "asdf", "--req-array", "qwerty")
+		i.Environ.Set(serpent.CompletionModeEnv, "1")
+		io := fakeIO(i)
+		err := i.Run()
+		require.NoError(t, err)
+		require.Equal(t, "--req-array\n--req-enum\n", io.Stdout.String())
+	})
+
+	t.Run("FlagShorthand", func(t *testing.T) {
+		t.Parallel()
+		i := cmd().Invoke("required-flag", "-b", "-s", "foo bar", "-a", "asdf")
 		i.Environ.Set(serpent.CompletionModeEnv, "1")
 		io := fakeIO(i)
 		err := i.Run()
@@ -74,6 +94,7 @@ func TestCompletion(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "--req-enum=foo\n--req-enum=bar\n--req-enum=qux\n", io.Stdout.String())
 	})
+
 }
 
 func TestFileCompletion(t *testing.T) {

example/completetest/main.go

@@ -26,12 +26,11 @@ func InstallCommand() *serpent.Command {
 		},
 		Options: serpent.OptionSet{
 			{
-				Flag:              "shell",
-				FlagShorthand:     "s",
-				Default:           defaultShell,
-				Description:       "The shell to generate a completion script for.",
-				Value:             completion.ShellOptions(&shell),
-				CompletionHandler: completion.ShellHandler(),
+				Flag:          "shell",
+				FlagShorthand: "s",
+				Default:       defaultShell,
+				Description:   "The shell to generate a completion script for.",
+				Value:         completion.ShellOptions(&shell),
 			},
 		},
 	}
@@ -99,11 +98,10 @@ func main() {
 						Description: "Print the file.",
 					},
 					{
-						Name:              "type",
-						Value:             serpent.EnumOf(&fileType, "binary", "text"),
-						Flag:              "type",
-						Description:       "The type of file.",
-						CompletionHandler: completion.EnumHandler("binary", "text"),
+						Name:        "type",
+						Value:       serpent.EnumOf(&fileType, "binary", "text"),
+						Flag:        "type",
+						Description: "The type of file.",
 					},
 					{
 						Name:        "extra",

option.go

@@ -349,7 +349,7 @@ func (optSet OptionSet) ByName(name string) *Option {
 func (optSet OptionSet) ByFlag(flag string) *Option {
 	for i := range optSet {
 		opt := &optSet[i]
-		if opt.Flag == flag {
+		if opt.Flag == flag || opt.FlagShorthand == flag {
 			return opt
 		}
 	}