From 33640667617ef83587c43a720ebac3e8032ec9ef Mon Sep 17 00:00:00 2001
From: JonasG <jonas.grgt@gmail.com>
Date: Sun, 17 Nov 2024 08:25:49 +0100
Subject: [PATCH] feat: add option to not quit form after submission

---
 form.go                       | 16 ++++--
 huh_test.go                   | 98 +++++++++++++++++++++++++++++++++++
 internal/selector/selector.go |  2 +-
 3 files changed, 112 insertions(+), 4 deletions(-)

diff --git a/form.go b/form.go
index 02d34a1d..cbbb5c04 100644
--- a/form.go
+++ b/form.go
@@ -76,8 +76,12 @@ type Form struct {
 	// to be more accessible to screen readers.
 	accessible bool
 
-	quitting bool
-	aborted  bool
+	// whether to quit the form after Submitting or not,
+	// defaults to true
+	QuitAfterSubmit bool
+	prevShowHelp    bool
+	quitting        bool
+	aborted         bool
 
 	// options
 	width      int
@@ -105,6 +109,8 @@ func NewForm(groups ...*Group) *Form {
 		teaOptions: []tea.ProgramOption{
 			tea.WithOutput(os.Stderr),
 		},
+		QuitAfterSubmit: true,
+		prevShowHelp:    true,
 	}
 
 	// NB: If dynamic forms come into play this will need to be applied when
@@ -494,6 +500,8 @@ func (f *Form) Init() tea.Cmd {
 		cmds = append(cmds, nextGroup)
 	}
 
+	f.selector.Selected().showHelp = f.prevShowHelp
+
 	return tea.Batch(cmds...)
 }
 
@@ -551,6 +559,8 @@ func (f *Form) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 
 		if f.selector.OnLast() {
+			f.prevShowHelp = f.selector.Selected().showHelp
+			f.selector.Selected().showHelp = false
 			return submit()
 		}
 
@@ -607,7 +617,7 @@ func (f *Form) isGroupHidden(group *Group) bool {
 
 // View renders the form.
 func (f *Form) View() string {
-	if f.quitting {
+	if f.quitting && f.QuitAfterSubmit {
 		return ""
 	}
 
diff --git a/huh_test.go b/huh_test.go
index 57288b3b..5bff37ed 100644
--- a/huh_test.go
+++ b/huh_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/charmbracelet/bubbles/cursor"
 	"io"
 	"regexp"
 	"strings"
@@ -275,6 +276,86 @@ func TestForm(t *testing.T) {
 	// TODO: Finish and submit form.
 }
 
+func TestFormNotQuiting(t *testing.T) {
+	t.Run("Quit form after submission", func(t *testing.T) {
+		f := NewForm(NewGroup(NewInput()))
+		f.Update(f.Init())
+
+		f = submitForm(f)
+
+		view := ansi.Strip(f.View())
+
+		if strings.Contains(view, ">") {
+			t.Log(pretty.Render(view))
+			t.Error("Expected form to have quit but was visible.")
+		}
+	})
+
+	t.Run("Form stays alive after submission", func(t *testing.T) {
+		f := NewForm(NewGroup(NewInput()))
+		f.QuitAfterSubmit = false
+		f.Update(f.Init())
+
+		submitForm(f)
+
+		view := ansi.Strip(f.View())
+
+		if !strings.Contains(view, ">") {
+			t.Log(pretty.Render(view))
+			t.Error("Expected form to be visible but has quit.")
+		}
+	})
+
+	t.Run("Grouped form stays alive after submission", func(t *testing.T) {
+		f := NewForm(NewGroup(NewInput()), NewGroup(NewInput()))
+		f.QuitAfterSubmit = false
+		f.Update(f.Init())
+
+		submitForm(f)
+		submitForm(f)
+
+		view := ansi.Strip(f.View())
+
+		if !strings.Contains(view, ">") {
+			t.Log(pretty.Render(view))
+			t.Error("Expected grouped form to be visible but has quit.")
+		}
+	})
+
+	t.Run("Hide group help when frozen", func(t *testing.T) {
+		f := NewForm(NewGroup(NewInput()))
+		f.QuitAfterSubmit = false
+		f.Update(f.Init())
+
+		submitForm(f)
+		submitForm(f)
+
+		view := ansi.Strip(f.View())
+
+		if strings.Contains(view, "enter submit") {
+			t.Log(pretty.Render(view))
+			t.Error("Expected help to be hidden but was visible.")
+		}
+	})
+
+	t.Run("Reset previous group help state after re-initialization", func(t *testing.T) {
+		f := NewForm(NewGroup(NewInput()).WithShowHelp(true))
+		f.QuitAfterSubmit = false
+		f.Update(f.Init())
+
+		submitForm(f)
+		submitForm(f)
+		f.Init()
+
+		view := ansi.Strip(f.View())
+
+		if !strings.Contains(view, "enter submit") {
+			t.Log(pretty.Render(view))
+			t.Error("Expected help visible but was not.")
+		}
+	})
+}
+
 func TestInput(t *testing.T) {
 	field := NewInput()
 	f := NewForm(NewGroup(field))
@@ -925,6 +1006,23 @@ func formProgram() *Form {
 		WithAccessible(false)
 }
 
+func submitForm(f *Form) *Form {
+	m, cmd := f.Update(tea.KeyMsg{Type: tea.KeyEnter})
+	f = m.(*Form)
+	for {
+		if cmd != nil {
+			msg := cmd()
+			_, cmd = f.Update(msg)
+			if _, ok := msg.(cursor.BlinkMsg); ok {
+				break
+			}
+		} else {
+			break
+		}
+	}
+	return f
+}
+
 func batchUpdate(m tea.Model, cmd tea.Cmd) tea.Model {
 	if cmd == nil {
 		return m
diff --git a/internal/selector/selector.go b/internal/selector/selector.go
index d0d41284..83961c54 100644
--- a/internal/selector/selector.go
+++ b/internal/selector/selector.go
@@ -52,7 +52,7 @@ func (s *Selector[T]) Index() int {
 	return s.index
 }
 
-// Totoal returns the total number of items.
+// Total returns the total number of items.
 func (s *Selector[T]) Total() int {
 	return len(s.items)
 }