Support suggestions generated from command in custom commands

This changes the interface a bit but it was only added earlier today so I doubt anybody is dependent on it yet.

I'm also updating the docs.

6 files changed, 136 insertions, 27 deletions

diff --git a/pkg/config/user_config.go b/pkg/config/user_config.go
index 27ae17953..9f46948b5 100644
--- a/pkg/config/user_config.go
+++ b/pkg/config/user_config.go
@@ -358,16 +358,14 @@ type CustomCommand struct {
 }
 
 type CustomCommandPrompt struct {
-	Key string `yaml:"key"`
-
 	// one of 'input', 'menu', 'confirm', or 'menuFromCommand'
-	Type string `yaml:"type"`
-
+	Type  string `yaml:"type"`
+	Key   string `yaml:"key"`
 	Title string `yaml:"title"`
 
-	// this only apply to input prompts
-	InitialValue      string `yaml:"initialValue"`
-	SuggestionsPreset string `yaml:"suggestionsPreset"`
+	// these only apply to input prompts
+	InitialValue string                   `yaml:"initialValue"`
+	Suggestions  CustomCommandSuggestions `yaml:"suggestions"`
 
 	// this only applies to confirm prompts
 	Body string `yaml:"body"`
@@ -382,6 +380,11 @@ type CustomCommandPrompt struct {
 	LabelFormat string `yaml:"labelFormat"`
 }
 
+type CustomCommandSuggestions struct {
+	Preset  string `yaml:"preset"`
+	Command string `yaml:"command"`
+}
+
 type CustomCommandMenuOption struct {
 	Name        string `yaml:"name"`
 	Description string `yaml:"description"`
diff --git a/pkg/gui/services/custom_commands/handler_creator.go b/pkg/gui/services/custom_commands/handler_creator.go
index 2d13a6f55..8e6700f72 100644
--- a/pkg/gui/services/custom_commands/handler_creator.go
+++ b/pkg/gui/services/custom_commands/handler_creator.go
@@ -11,6 +11,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
+	"github.com/samber/lo"
 )
 
 // takes a custom command and returns a function that will be called when the corresponding user-defined keybinding is pressed
@@ -108,13 +109,9 @@ func (self *HandlerCreator) call(customCommand config.CustomCommand) func() erro
 }
 
 func (self *HandlerCreator) inputPrompt(prompt *config.CustomCommandPrompt, wrappedF func(string) error) error {
-	var findSuggestionsFn func(string) []*types.Suggestion
-	if prompt.SuggestionsPreset != "" {
-		var err error
-		findSuggestionsFn, err = self.getPresetSuggestionsFn(prompt.SuggestionsPreset)
-		if err != nil {
-			return err
-		}
+	findSuggestionsFn, err := self.generateFindSuggestionsFunc(prompt)
+	if err != nil {
+		return self.c.Error(err)
 	}
 
 	return self.c.Prompt(types.PromptOpts{
@@ -127,6 +124,41 @@ func (self *HandlerCreator) inputPrompt(prompt *config.CustomCommandPrompt, wrap
 	})
 }
 
+func (self *HandlerCreator) generateFindSuggestionsFunc(prompt *config.CustomCommandPrompt) (func(string) []*types.Suggestion, error) {
+	if prompt.Suggestions.Preset != "" && prompt.Suggestions.Command != "" {
+		return nil, fmt.Errorf(
+			fmt.Sprintf(
+				"Custom command prompt cannot have both a preset and a command for suggestions. Preset: '%s', Command: '%s'",
+				prompt.Suggestions.Preset,
+				prompt.Suggestions.Command,
+			),
+		)
+	} else if prompt.Suggestions.Preset != "" {
+		return self.getPresetSuggestionsFn(prompt.Suggestions.Preset)
+	} else if prompt.Suggestions.Command != "" {
+		return self.getCommandSuggestionsFn(prompt.Suggestions.Command)
+	}
+
+	return nil, nil
+}
+
+func (self *HandlerCreator) getCommandSuggestionsFn(command string) (func(string) []*types.Suggestion, error) {
+	lines := []*types.Suggestion{}
+	err := self.c.OS().Cmd.NewShell(command).RunAndProcessLines(func(line string) (bool, error) {
+		lines = append(lines, &types.Suggestion{Value: line, Label: line})
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return func(currentWord string) []*types.Suggestion {
+		return lo.Filter(lines, func(suggestion *types.Suggestion, _ int) bool {
+			return strings.Contains(strings.ToLower(suggestion.Value), strings.ToLower(currentWord))
+		})
+	}, nil
+}
+
 func (self *HandlerCreator) getPresetSuggestionsFn(preset string) (func(string) []*types.Suggestion, error) {
 	switch preset {
 	case "files":
@@ -144,6 +176,14 @@ func (self *HandlerCreator) getPresetSuggestionsFn(preset string) (func(string)
 	}
 }
 
+func (self *HandlerCreator) confirmPrompt(prompt *config.CustomCommandPrompt, handleConfirm func() error) error {
+	return self.c.Confirm(types.ConfirmOpts{
+		Title:         prompt.Title,
+		Prompt:        prompt.Body,
+		HandleConfirm: handleConfirm,
+	})
+}
+
 func (self *HandlerCreator) menuPrompt(prompt *config.CustomCommandPrompt, wrappedF func(string) error) error {
 	menuItems := slices.Map(prompt.Options, func(option config.CustomCommandMenuOption) *types.MenuItem {
 		return &types.MenuItem{
@@ -157,14 +197,6 @@ func (self *HandlerCreator) menuPrompt(prompt *config.CustomCommandPrompt, wrapp
 	return self.c.Menu(types.CreateMenuOptions{Title: prompt.Title, Items: menuItems})
 }
 
-func (self *HandlerCreator) confirmPrompt(prompt *config.CustomCommandPrompt, handleConfirm func() error) error {
-	return self.c.Confirm(types.ConfirmOpts{
-		Title:         prompt.Title,
-		Prompt:        prompt.Body,
-		HandleConfirm: handleConfirm,
-	})
-}
-
 func (self *HandlerCreator) menuPromptFromCommand(prompt *config.CustomCommandPrompt, wrappedF func(string) error) error {
 	// Run and save output
 	message, err := self.c.Git().Custom.RunWithOutput(prompt.Command)
diff --git a/pkg/gui/services/custom_commands/resolver.go b/pkg/gui/services/custom_commands/resolver.go
index 119581bfc..67dd7be2d 100644
--- a/pkg/gui/services/custom_commands/resolver.go
+++ b/pkg/gui/services/custom_commands/resolver.go
@@ -34,7 +34,12 @@ func (self *Resolver) resolvePrompt(
 		return nil, err
 	}
 
-	result.SuggestionsPreset, err = resolveTemplate(prompt.SuggestionsPreset)
+	result.Suggestions.Preset, err = resolveTemplate(prompt.Suggestions.Preset)
+	if err != nil {
+		return nil, err
+	}
+
+	result.Suggestions.Command, err = resolveTemplate(prompt.Suggestions.Command)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/integration/tests/custom_commands/suggestions_command.go b/pkg/integration/tests/custom_commands/suggestions_command.go
new file mode 100644
index 000000000..f54ec1bae
--- /dev/null
+++ b/pkg/integration/tests/custom_commands/suggestions_command.go
@@ -0,0 +1,66 @@
+package custom_commands
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var SuggestionsCommand = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Using a custom command that uses a suggestions command in a prompt step",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("branch-one")
+		shell.EmptyCommit("blah")
+		shell.NewBranch("branch-two")
+		shell.EmptyCommit("blah")
+		shell.NewBranch("branch-three")
+		shell.EmptyCommit("blah")
+		shell.NewBranch("branch-four")
+		shell.EmptyCommit("blah")
+	},
+	SetupConfig: func(cfg *config.AppConfig) {
+		cfg.UserConfig.CustomCommands = []config.CustomCommand{
+			{
+				Key:     "a",
+				Context: "localBranches",
+				Command: `git checkout {{.Form.Branch}}`,
+				Prompts: []config.CustomCommandPrompt{
+					{
+						Key:   "Branch",
+						Type:  "input",
+						Title: "Enter a branch name",
+						Suggestions: config.CustomCommandSuggestions{
+							Command: "git branch --format='%(refname:short)'",
+						},
+					},
+				},
+			},
+		}
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Branches().
+			Focus().
+			Lines(
+				Contains("branch-four").IsSelected(),
+				Contains("branch-three"),
+				Contains("branch-two"),
+				Contains("branch-one"),
+			).
+			Press("a")
+
+		t.ExpectPopup().Prompt().
+			Title(Equals("Enter a branch name")).
+			Type("three").
+			SuggestionLines(Contains("branch-three")).
+			ConfirmFirstSuggestion()
+
+		t.Views().Branches().
+			Lines(
+				Contains("branch-three").IsSelected(),
+				Contains("branch-four"),
+				Contains("branch-two"),
+				Contains("branch-one"),
+			)
+	},
+})
diff --git a/pkg/integration/tests/custom_commands/suggestions_preset.go b/pkg/integration/tests/custom_commands/suggestions_preset.go
index 894e3b1fe..40c888625 100644
--- a/pkg/integration/tests/custom_commands/suggestions_preset.go
+++ b/pkg/integration/tests/custom_commands/suggestions_preset.go
@@ -27,10 +27,12 @@ var SuggestionsPreset = NewIntegrationTest(NewIntegrationTestArgs{
 				Command: `git checkout {{.Form.Branch}}`,
 				Prompts: []config.CustomCommandPrompt{
 					{
-						Key:               "Branch",
-						Type:              "input",
-						Title:             "Enter a branch name",
-						SuggestionsPreset: "branches",
+						Key:   "Branch",
+						Type:  "input",
+						Title: "Enter a branch name",
+						Suggestions: config.CustomCommandSuggestions{
+							Preset: "branches",
+						},
 					},
 				},
 			},
diff --git a/pkg/integration/tests/test_list.go b/pkg/integration/tests/test_list.go
index e0800f7bd..50377f101 100644
--- a/pkg/integration/tests/test_list.go
+++ b/pkg/integration/tests/test_list.go
@@ -77,6 +77,7 @@ var tests = []*components.IntegrationTest{
 	custom_commands.MenuFromCommandsOutput,
 	custom_commands.MultiplePrompts,
 	custom_commands.OmitFromHistory,
+	custom_commands.SuggestionsCommand,
 	custom_commands.SuggestionsPreset,
 	diff.Diff,
 	diff.DiffAndApplyPatch,