Repository: charmbracelet/bubbletea
Pull request #715 (2023-04-07 15:08:27)

fix(output): reuse termenv output
If the passed io.Writer is a termenv.Output, use it instead of creating a new termenv.Output.


Patch:
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
index 6d77155..0e7a6cb 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -12,7 +12,7 @@ A clear and concise description of what the bug is.
 
 **Setup**
 Please complete the following information along with version numbers, if applicable.
- - OS [e.g. Ubuntu, Mac OS]
+ - OS [e.g. Ubuntu, macOS]
  - Shell [e.g. zsh, fish]
  - Terminal Emulator [e.g. kitty, iterm]
  - Terminal Multiplexer [e.g. tmux]
@@ -25,7 +25,7 @@ Steps to reproduce the behavior:
 4. See error
 
 **Source Code**
-Please include source code if needed to reproduce the behaviour. 
+Please include source code if needed to reproduce the behavior. 
 
 **Expected behavior**
 A clear and concise description of what you expected to happen.
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..a060722 100644
--- a/options.go
+++ b/options.go
@@ -28,7 +28,11 @@ func WithContext(ctx context.Context) ProgramOption {
 // won't need to use this.
 func WithOutput(output io.Writer) ProgramOption {
 	return func(p *Program) {
-		p.output = termenv.NewOutput(output, termenv.WithColorCache(true))
+		if o, ok := output.(*termenv.Output); ok {
+			p.output = o
+		} else {
+			p.output = termenv.NewOutput(output, termenv.WithColorCache(true))
+		}
 	}
 }
 
@@ -151,3 +155,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/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:


Test patch:
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_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