review p2

Author: ethanndickson

command.go

@@ -407,31 +407,8 @@ func (inv *Invocation) run(state *runState) error {
 	// Outputted completions are not filtered based on the word under the cursor, as every shell we support does this already.
 	// We only look at the current word to figure out handler to run, or what directory to inspect.
 	if inv.IsCompletionMode() {
-		prev, cur := inv.curWords()
-		inv.CurWord = cur
-		// If the current word is a flag set using `=`, use it's handler
-		if strings.HasPrefix(cur, "--") && strings.Contains(cur, "=") {
-			if inv.equalsFlagHandler(cur) {
-				return nil
-			}
-		}
-		// If the previous word is a flag, then we're writing it's value
-		// and we should check it's handler
-		if strings.HasPrefix(prev, "--") {
-			if inv.flagHandler(prev) {
-				return nil
-			}
-		}
-		// If the current word is the command, auto-complete it so the shell moves the cursor
-		if inv.Command.Name() == inv.CurWord {
-			fmt.Fprintf(inv.Stdout, "%s\n", inv.Command.Name())
-			return nil
-		}
-		if inv.Command.CompletionHandler == nil {
-			inv.Command.CompletionHandler = DefaultCompletionHandler
-		}
-		for _, e := range inv.Command.CompletionHandler(inv) {
-			fmt.Fprintf(inv.Stdout, "%s\n", e)
+		for _, e := range inv.doCompletions() {
+			fmt.Fprintln(inv.Stdout, e)
 		}
 		return nil
 	}
@@ -613,11 +590,42 @@ func (inv *Invocation) with(fn func(*Invocation)) *Invocation {
 	return &i2
 }
 
-func (inv *Invocation) flagHandler(word string) bool {
-	return inv.doFlagCompletion("", word)
+func (inv *Invocation) doCompletions() []string {
+	prev, cur := inv.curWords()
+	inv.CurWord = cur
+	// If the current word is a flag set using `=`, use it's handler
+	if strings.HasPrefix(cur, "--") && strings.Contains(cur, "=") {
+		if out := inv.equalsFlagCompletions(cur); out != nil {
+			return out
+		}
+	}
+	// If the previous word is a flag, then we're writing it's value
+	// and we should check it's handler
+	if strings.HasPrefix(prev, "--") {
+		if out := inv.flagCompletions(prev); out != nil {
+			return out
+		}
+	}
+	// If the current word is the command, auto-complete it so the shell moves the cursor
+	if inv.Command.Name() == inv.CurWord {
+		return []string{inv.Command.Name()}
+	}
+	var completions []string
+
+	if inv.Command.CompletionHandler != nil {
+		completions = append(completions, inv.Command.CompletionHandler(inv)...)
+	}
+
+	completions = append(completions, DefaultCompletionHandler(inv)...)
+
+	return completions
 }
 
-func (inv *Invocation) equalsFlagHandler(word string) bool {
+func (inv *Invocation) flagCompletions(word string) []string {
+	return inv.doFlagCompletions("", word)
+}
+
+func (inv *Invocation) equalsFlagCompletions(word string) []string {
 	words := strings.Split(word, "=")
 	word = words[0]
 	if len(words) > 1 {
@@ -626,29 +634,32 @@ func (inv *Invocation) equalsFlagHandler(word string) bool {
 		inv.CurWord = ""
 	}
 	prefix := word + "="
-	return inv.doFlagCompletion(prefix, word)
+	return inv.doFlagCompletions(prefix, word)
 }
 
-func (inv *Invocation) doFlagCompletion(prefix, word string) bool {
+func (inv *Invocation) doFlagCompletions(prefix, word string) []string {
 	opt := inv.Command.Options.ByFlag(word[2:])
 	if opt == nil {
-		return false
+		return nil
 	}
 	if opt.CompletionHandler != nil {
 		completions := opt.CompletionHandler(inv)
+		out := make([]string, 0, len(completions))
 		for _, completion := range completions {
-			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, completion)
+			out = append(out, fmt.Sprintf("%s%s", prefix, completion))
 		}
-		return true
+		return out
 	}
 	val, ok := opt.Value.(*Enum)
 	if ok {
-		for _, choice := range val.Choices {
-			fmt.Fprintf(inv.Stdout, "%s%s\n", prefix, choice)
+		completions := val.Choices
+		out := make([]string, 0, len(completions))
+		for _, choice := range completions {
+			out = append(out, fmt.Sprintf("%s%s", prefix, choice))
 		}
-		return true
+		return out
 	}
-	return false
+	return nil
 }
 
 // MiddlewareFunc returns the next handler in the chain,

command_test.go

@@ -161,6 +161,9 @@ func sampleCommand(t *testing.T) *serpent.Command {
 						CompletionHandler: completion.FileListHandler(nil),
 					},
 				},
+				CompletionHandler: func(i *serpent.Invocation) []string {
+					return []string{"doesntexist.go"}
+				},
 			},
 		},
 	}

completion/handlers.go

@@ -36,7 +36,8 @@ func FileListHandler(filter func(info os.FileInfo) bool) serpent.CompletionHandl
 }
 
 func listFiles(word string, filter func(info os.FileInfo) bool) []string {
-	out := make([]string, 0, 32)
+	// Avoid reallocating for each of the first few files we see.
+	out := make([]string, 0, 16)
 
 	dir, _ := filepath.Split(word)
 	if dir == "" {

completion_test.go

@@ -55,6 +55,16 @@ func TestCompletion(t *testing.T) {
 		require.Equal(t, "--req-array\n--req-bool\n--req-enum\n--req-string\n", io.Stdout.String())
 	})
 
+	t.Run("ListFlagsAfterArg", func(t *testing.T) {
+		t.Parallel()
+		i := cmd().Invoke("altfile", "")
+		i.Environ.Set(serpent.CompletionModeEnv, "1")
+		io := fakeIO(i)
+		err := i.Run()
+		require.NoError(t, err)
+		require.Equal(t, "doesntexist.go\n--extra\n", io.Stdout.String())
+	})
+
 	t.Run("FlagExhaustive", func(t *testing.T) {
 		t.Parallel()
 		i := cmd().Invoke("required-flag", "--req-bool", "--req-string", "foo bar", "--req-array", "asdf", "--req-array", "qwerty")