From 9545e7cee2102a5f98aa33449f44b618339cc713 Mon Sep 17 00:00:00 2001 From: Christian Muehlhaeuser Date: Thu, 13 Oct 2022 07:29:07 +0200 Subject: [PATCH] feat: add generic event filter 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 | 15 ++++++++++----- options.go | 36 ++++++++++++++++++++--------------- options_test.go | 8 ++++---- tea.go | 26 +++++++++---------------- tea_test.go | 19 ++++++++++-------- 5 files changed, 55 insertions(+), 49 deletions(-) diff --git a/examples/prevent-quit/main.go b/examples/prevent-quit/main.go index 1916b0ca70..1339393937 100644 --- a/examples/prevent-quit/main.go +++ b/examples/prevent-quit/main.go @@ -1,6 +1,6 @@ package main -// A program demonstrating how to use the WithOnQuit option to intercept quit events +// A program demonstrating how to use the WithFilter option to intercept events. import ( "fmt" @@ -20,19 +20,24 @@ var ( ) func main() { - p := tea.NewProgram(initialModel(), tea.WithOnQuit(onQuit)) + p := tea.NewProgram(initialModel(), tea.WithFilter(filter)) if _, err := p.Run(); err != nil { log.Fatal(err) } } -func onQuit(teaModel tea.Model) tea.QuitBehavior { +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 tea.PreventShutdown + return nil } - return tea.Shutdown + + return msg } type model struct { diff --git a/options.go b/options.go index a76720eabb..4e480b40b2 100644 --- a/options.go +++ b/options.go @@ -142,31 +142,37 @@ func WithANSICompressor() ProgramOption { } } -// 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. +// 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 onQuit(m tea.Model) tea.QuitBehavior { -// model := m.(myModel) -// if model.hasChanges { -// return tea.PreventShutdown -// } -// return tea.Shutdown +// 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.WithOnQuit(onQuit)); +// p := tea.NewProgram(Model{}, tea.WithFilter(filter)); // // if _,err := p.Run(); err != nil { // fmt.Println("Error running program:", err) // os.Exit(1) // } -func WithOnQuit(onQuit func(Model) QuitBehavior) ProgramOption { +func WithFilter(filter func(Model, Msg) Msg) ProgramOption { return func(p *Program) { - p.onQuit = onQuit + p.filter = filter } } diff --git a/options_test.go b/options_test.go index 1c406494a9..71a3c6c6d0 100644 --- a/options_test.go +++ b/options_test.go @@ -35,10 +35,10 @@ 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("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") } }) diff --git a/tea.go b/tea.go index 1da6bd9d5d..af9f61ae27 100644 --- a/tea.go +++ b/tea.go @@ -124,23 +124,10 @@ type Program struct { // below. windowsStdin *os.File //nolint:golint,structcheck,unused - onQuit func(Model) QuitBehavior + filter func(Model, Msg) Msg } -// 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{} } @@ -286,12 +273,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: - if p.onQuit != nil && p.onQuit(model) == PreventShutdown { - break - } return model, nil case clearScreenMsg: diff --git a/tea_test.go b/tea_test.go index 51ec526b93..f4a41ff523 100644 --- a/tea_test.go +++ b/tea_test.go @@ -76,13 +76,13 @@ func TestTeaQuit(t *testing.T) { } } -func TestTeaWithOnQuit(t *testing.T) { - testTeaWithOnQuit(t, 0) - testTeaWithOnQuit(t, 1) - testTeaWithOnQuit(t, 2) +func TestTeaWithFilter(t *testing.T) { + testTeaWithFilter(t, 0) + testTeaWithFilter(t, 1) + testTeaWithFilter(t, 2) } -func testTeaWithOnQuit(t *testing.T, preventCount uint32) { +func testTeaWithFilter(t *testing.T, preventCount uint32) { var buf bytes.Buffer var in bytes.Buffer @@ -91,12 +91,15 @@ func testTeaWithOnQuit(t *testing.T, preventCount uint32) { p := NewProgram(m, WithInput(&in), WithOutput(&buf), - WithOnQuit(func(Model) QuitBehavior { + WithFilter(func(_ Model, msg Msg) Msg { + if _, ok := msg.(QuitMsg); !ok { + return msg + } if shutdowns < preventCount { atomic.AddUint32(&shutdowns, 1) - return PreventShutdown + return nil } - return Shutdown + return msg })) go func() {