From c56884c0e247cafa6bfa3e222aa512fe6d6a292e Mon Sep 17 00:00:00 2001 From: Christian Muehlhaeuser Date: Mon, 17 Apr 2023 22:02:55 +0200 Subject: [PATCH] feat: add generic event filter (#536) `WithFilter` lets you supply an event filter that will be invoked before Bubble Tea processes a `tea.Msg`. The event filter can return any `tea.Msg` which will then get handled by Bubble Tea instead of the original event. If the event filter returns nil, the event will be ignored and Bubble Tea will not process it. As an example, this could be used to prevent a program from shutting down if there are unsaved changes. Based on the fantastic work by @aschey and supersedes #521. Resolves #472. --- examples/prevent-quit/main.go | 154 ++++++++++++++++++++++++++++++++++ options.go | 35 ++++++++ options_test.go | 7 ++ tea.go | 22 +++-- tea_test.go | 41 +++++++++ 5 files changed, 253 insertions(+), 6 deletions(-) create mode 100644 examples/prevent-quit/main.go diff --git a/examples/prevent-quit/main.go b/examples/prevent-quit/main.go new file mode 100644 index 0000000..1339393 --- /dev/null +++ b/examples/prevent-quit/main.go @@ -0,0 +1,154 @@ +package main + +// A program demonstrating how to use the WithFilter option to intercept 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.WithFilter(filter)) + + if _, err := p.Run(); err != nil { + log.Fatal(err) + } +} + +func filter(teaModel tea.Model, msg tea.Msg) tea.Msg { + if _, ok := msg.(tea.QuitMsg); !ok { + return msg + } + + m := teaModel.(model) + if m.hasChanges { + return nil + } + + return msg +} + +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" +} diff --git a/options.go b/options.go index 9945924..17d05ba 100644 --- a/options.go +++ b/options.go @@ -151,3 +151,38 @@ func WithANSICompressor() ProgramOption { p.startupOptions |= withANSICompressor } } + +// WithFilter supplies an event filter that will be invoked before Bubble Tea +// processes a tea.Msg. The event filter can return any tea.Msg which will then +// get handled by Bubble Tea instead of the original event. If the event filter +// returns nil, the event will be ignored and Bubble Tea will not process it. +// +// As an example, this could be used to prevent a program from shutting down if +// there are unsaved changes. +// +// Example: +// +// func filter(m tea.Model, msg tea.Msg) tea.Msg { +// if _, ok := msg.(tea.QuitMsg); !ok { +// return msg +// } +// +// model := m.(myModel) +// if model.hasChanges { +// return nil +// } +// +// return msg +// } +// +// p := tea.NewProgram(Model{}, tea.WithFilter(filter)); +// +// if _,err := p.Run(); err != nil { +// fmt.Println("Error running program:", err) +// os.Exit(1) +// } +func WithFilter(filter func(Model, Msg) Msg) ProgramOption { + return func(p *Program) { + p.filter = filter + } +} diff --git a/options_test.go b/options_test.go index 7e08f58..71a3c6c 100644 --- a/options_test.go +++ b/options_test.go @@ -35,6 +35,13 @@ func TestOptions(t *testing.T) { } }) + t.Run("filter", func(t *testing.T) { + p := NewProgram(nil, WithFilter(func(_ Model, msg Msg) Msg { return msg })) + if p.filter == nil { + t.Errorf("expected filter to be set") + } + }) + t.Run("startup options", func(t *testing.T) { exercise := func(t *testing.T, opt ProgramOption, expect startupOptions) { p := NewProgram(nil, opt) diff --git a/tea.go b/tea.go index 76c9b3b..f5ae4e9 100644 --- a/tea.go +++ b/tea.go @@ -123,16 +123,18 @@ type Program struct { // as this value only comes into play on Windows, hence the ignore comment // below. windowsStdin *os.File //nolint:golint,structcheck,unused + + filter func(Model, Msg) Msg } // Quit is a special command that tells the Bubble Tea program to exit. 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 { @@ -194,7 +196,7 @@ func (p *Program) handleSignals() chan struct{} { case <-sig: if !p.ignoreSignals { - p.msgs <- quitMsg{} + p.msgs <- QuitMsg{} return } } @@ -267,9 +269,17 @@ func (p *Program) eventLoop(model Model, cmds chan Cmd) (Model, error) { return model, err case msg := <-p.msgs: + // Filter messages. + if p.filter != nil { + msg = p.filter(model, msg) + } + if msg == nil { + continue + } + // Handle special internal messages. switch msg := msg.(type) { - case quitMsg: + case QuitMsg: return model, nil case clearScreenMsg: diff --git a/tea_test.go b/tea_test.go index 748b4cb..d65f0d7 100644 --- a/tea_test.go +++ b/tea_test.go @@ -77,6 +77,47 @@ func TestTeaQuit(t *testing.T) { } } +func TestTeaWithFilter(t *testing.T) { + testTeaWithFilter(t, 0) + testTeaWithFilter(t, 1) + testTeaWithFilter(t, 2) +} + +func testTeaWithFilter(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), + WithFilter(func(_ Model, msg Msg) Msg { + if _, ok := msg.(QuitMsg); !ok { + return msg + } + if shutdowns < preventCount { + atomic.AddUint32(&shutdowns, 1) + return nil + } + return msg + })) + + 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