feat: add quit event handler

examples/prevent-quit/main.go

@@ -0,0 +1,149 @@
+package main
+
+// A program demonstrating how to use the WithOnQuit option to intercept quit events
+
+import (
+ "fmt"
+ "log"
+
+ "github.com/charmbracelet/bubbles/help"
+ "github.com/charmbracelet/bubbles/key"
+ "github.com/charmbracelet/bubbles/textarea"
+ tea "github.com/charmbracelet/bubbletea"
+ "github.com/charmbracelet/lipgloss"
+)
+
+var (
+ choiceStyle = lipgloss.NewStyle().PaddingLeft(1).Foreground(lipgloss.Color("241"))
+ saveTextStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("170"))
+ quitViewStyle = lipgloss.NewStyle().Padding(1).Border(lipgloss.RoundedBorder()).BorderForeground(lipgloss.Color("170"))
+)
+
+func main() {
+ p := tea.NewProgram(initialModel(), tea.WithOnQuit(onQuit))
+
+ if _, err := p.Run(); err != nil {
+ log.Fatal(err)
+ }
+}
+
+func onQuit(teaModel tea.Model) tea.QuitBehavior {
+ m := teaModel.(model)
+ if m.hasChanges {
+ return tea.PreventShutdown
+ }
+ return tea.Shutdown
+}
+
+type model struct {
+ textarea textarea.Model
+ help help.Model
+ keymap keymap
+ saveText string
+ hasChanges bool
+ quitting bool
+}
+
+type keymap struct {
+ save key.Binding
+ quit key.Binding
+}
+
+func initialModel() model {
+ ti := textarea.New()
+ ti.Placeholder = "Only the best words"
+ ti.Focus()
+
+ return model{
+ textarea: ti,
+ help: help.NewModel(),
+ keymap: keymap{
+ save: key.NewBinding(
+ key.WithKeys("ctrl+s"),
+ key.WithHelp("ctrl+s", "save"),
+ ),
+ quit: key.NewBinding(
+ key.WithKeys("esc", "ctrl+c"),
+ key.WithHelp("esc", "quit"),
+ ),
+ },
+ }
+}
+
+func (m model) Init() tea.Cmd {
+ return textarea.Blink
+}
+
+func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+ if m.quitting {
+ return m.updatePromptView(msg)
+ }
+
+ return m.updateTextView(msg)
+}
+
+func (m model) updateTextView(msg tea.Msg) (tea.Model, tea.Cmd) {
+ var cmds []tea.Cmd
+ var cmd tea.Cmd
+
+ switch msg := msg.(type) {
+ case tea.KeyMsg:
+ m.saveText = ""
+ switch {
+ case key.Matches(msg, m.keymap.save):
+ m.saveText = "Changes saved!"
+ m.hasChanges = false
+ case key.Matches(msg, m.keymap.quit):
+ m.quitting = true
+ return m, tea.Quit
+ case msg.Type == tea.KeyRunes:
+ m.saveText = ""
+ m.hasChanges = true
+ fallthrough
+ default:
+ if !m.textarea.Focused() {
+ cmd = m.textarea.Focus()
+ cmds = append(cmds, cmd)
+ }
+ }
+ }
+ m.textarea, cmd = m.textarea.Update(msg)
+ cmds = append(cmds, cmd)
+ return m, tea.Batch(cmds...)
+}
+
+func (m model) updatePromptView(msg tea.Msg) (tea.Model, tea.Cmd) {
+ switch msg := msg.(type) {
+ case tea.KeyMsg:
+ // For simplicity's sake, we'll treat any key besides "y" as "no"
+ if key.Matches(msg, m.keymap.quit) || msg.String() == "y" {
+ m.hasChanges = false
+ return m, tea.Quit
+ }
+ m.quitting = false
+ }
+
+ return m, nil
+}
+
+func (m model) View() string {
+ if m.quitting {
+ if m.hasChanges {
+ text := lipgloss.JoinHorizontal(lipgloss.Top, "You have unsaved changes. Quit without saving?", choiceStyle.Render("[yn]"))
+ return quitViewStyle.Render(text)
+ }
+ return "Very important, thank you\n"
+ }
+
+ helpView := m.help.ShortHelpView([]key.Binding{
+ m.keymap.save,
+ m.keymap.quit,
+ })
+
+ return fmt.Sprintf(
+ "\nType some important things.\n\n%s\n\n %s\n %s",
+ m.textarea.View(),
+ saveTextStyle.Render(m.saveText),
+ helpView,
+ ) + "\n\n"
+}

options.go

@@ -141,3 +141,32 @@ func WithANSICompressor() ProgramOption {
  p.startupOptions |= withANSICompressor
  }
 }
+
+// WithOnQuit supplies an event handler that will be invoked whenever Bubble
+// Tea receives a QuitMsg. The event handler can return tea.Shutdown to
+// instruct Bubble Tea to handle the QuitMsg normally and shut the program
+// down, or it can return tea.PreventShutdown to prevent the program from
+// shutting down and instead handle the QuitMsg like a normal message and
+// pass it along to the model's Update method.
+//
+// Example:
+//
+// func onQuit(m tea.Model) tea.QuitBehavior {
+// model := m.(myModel)
+// if model.hasChanges {
+// return tea.PreventShutdown
+// }
+// return tea.Shutdown
+// }
+//
+// p := tea.NewProgram(Model{}, tea.WithOnQuit(onQuit));
+//
+// if _,err := p.Run(); err != nil {
+// fmt.Println("Error running program:", err)
+// os.Exit(1)
+// }
+func WithOnQuit(onQuit func(Model) QuitBehavior) ProgramOption {
+ return func(p *Program) {
+ p.onQuit = onQuit
+ }
+}

options_test.go

@@ -35,6 +35,13 @@ func TestOptions(t *testing.T) {
  }
  })
 
+ t.Run("on quit", func(t *testing.T) {
+ p := NewProgram(nil, WithOnQuit(func(Model) QuitBehavior { return Shutdown }))
+ if p.onQuit == nil {
+ t.Errorf("expected onQuit to be set")
+ }
+ })
+
  t.Run("startup options", func(t *testing.T) {
  exercise := func(t *testing.T, opt ProgramOption, expect startupOptions) {
  p := NewProgram(nil, opt)

tea.go

@@ -123,16 +123,31 @@ type Program struct {
  // as this value only comes into play on Windows, hence the ignore comment
  // below.
  windowsStdin *os.File //nolint:golint,structcheck,unused
+
+ onQuit func(Model) QuitBehavior
 }
 
+// QuitBehavior defines how Bubble Tea handles QuitMsgs.
+type QuitBehavior int
+
+const (
+ // Shutdown instructs Bubble Tea to shut down the program normally when a
+ // QuitMsg is received.
+ Shutdown QuitBehavior = iota
+ // PreventShutdown instructs Bubble Tea to ignore the QuitMsg that it
+ // received and instead pass the message to the model's Update function.
+ PreventShutdown
+)
+
 // Quit is a special command that tells the Bubble Tea program to exit.
+// This behavior can be controlled using the WithOnQuit option.
 func Quit() Msg {
- return quitMsg{}
+ return QuitMsg{}
 }
 
-// quitMsg in an internal message signals that the program should quit. You can
-// send a quitMsg with Quit.
-type quitMsg struct{}
+// QuitMsg signals that the program should quit. You can send a QuitMsg with
+// Quit.
+type QuitMsg struct{}
 
 // NewProgram creates a new Program.
 func NewProgram(model Model, opts ...ProgramOption) *Program {
@@ -186,7 +201,7 @@ func (p *Program) handleSignals() chan struct{} {
 
  case <-sig:
  if !p.ignoreSignals {
- p.msgs <- quitMsg{}
+ p.msgs <- QuitMsg{}
  return
  }
  }
@@ -273,7 +288,10 @@ func (p *Program) eventLoop(model Model, cmds chan Cmd) (Model, error) {
  case msg := <-p.msgs:
  // Handle special internal messages.
  switch msg := msg.(type) {
- case quitMsg:
+ case QuitMsg:
+ if p.onQuit != nil && p.onQuit(model) == PreventShutdown {
+ break
+ }
  return model, nil
 
  case clearScreenMsg:

tea_test.go

@@ -76,6 +76,44 @@ func TestTeaQuit(t *testing.T) {
  }
 }
 
+func TestTeaWithOnQuit(t *testing.T) {
+ testTeaWithOnQuit(t, 0)
+ testTeaWithOnQuit(t, 1)
+ testTeaWithOnQuit(t, 2)
+}
+
+func testTeaWithOnQuit(t *testing.T, preventCount uint32) {
+ var buf bytes.Buffer
+ var in bytes.Buffer
+
+ m := &testModel{}
+ shutdowns := uint32(0)
+ p := NewProgram(m,
+ WithInput(&in),
+ WithOutput(&buf),
+ WithOnQuit(func(Model) QuitBehavior {
+ if shutdowns < preventCount {
+ atomic.AddUint32(&shutdowns, 1)
+ return PreventShutdown
+ }
+ return Shutdown
+ }))
+
+ go func() {
+ for atomic.LoadUint32(&shutdowns) <= preventCount {
+ time.Sleep(time.Millisecond)
+ p.Quit()
+ }
+ }()
+
+ if err := p.Start(); err != nil {
+ t.Fatal(err)
+ }
+ if shutdowns != preventCount {
+ t.Errorf("Expected %d prevented shutdowns, got %d", preventCount, shutdowns)
+ }
+}
+
 func TestTeaKill(t *testing.T) {
  var buf bytes.Buffer
  var in bytes.Buffer