feat: tea.WithContext ProgramOption to supply a context

WithContext lets you specify a context in which to run the Program.
This is useful if you want to cancel the execution from outside.
When a Program gets cancelled it will exit with an error
ErrProgramKilled.
This commit is contained in:
Christian Muehlhaeuser 2022-10-23 03:43:52 +02:00
parent 0f1ce7f2d9
commit e15bcb7e0e
3 changed files with 40 additions and 1 deletions

View File

@ -1,6 +1,7 @@
package tea package tea
import ( import (
"context"
"io" "io"
"github.com/muesli/termenv" "github.com/muesli/termenv"
@ -14,6 +15,15 @@ import (
// p := NewProgram(model, WithInput(someInput), WithOutput(someOutput)) // p := NewProgram(model, WithInput(someInput), WithOutput(someOutput))
type ProgramOption func(*Program) type ProgramOption func(*Program)
// WithContext lets you specify a context in which to run the Program. This is
// useful if you want to cancel the execution from outside. When a Program gets
// cancelled it will exit with an error ErrProgramKilled.
func WithContext(ctx context.Context) ProgramOption {
return func(p *Program) {
p.ctx = ctx
}
}
// WithOutput sets the output which, by default, is stdout. In most cases you // WithOutput sets the output which, by default, is stdout. In most cases you
// won't need to use this. // won't need to use this.
func WithOutput(output io.Writer) ProgramOption { func WithOutput(output io.Writer) ProgramOption {

8
tea.go
View File

@ -141,8 +141,14 @@ func NewProgram(model Model, opts ...ProgramOption) *Program {
msgs: make(chan Msg), msgs: make(chan Msg),
} }
// A context can be provided with a ProgramOption, but if none was provided
// we'll use the default background context.
if p.ctx == nil {
p.ctx = context.Background()
}
// Initialize context and teardown channel. // Initialize context and teardown channel.
p.ctx, p.cancel = context.WithCancel(context.Background()) p.ctx, p.cancel = context.WithCancel(p.ctx)
// Apply all options to the program. // Apply all options to the program.
for _, opt := range opts { for _, opt := range opts {

View File

@ -2,6 +2,7 @@ package tea
import ( import (
"bytes" "bytes"
"context"
"sync/atomic" "sync/atomic"
"testing" "testing"
"time" "time"
@ -97,6 +98,28 @@ func TestTeaKill(t *testing.T) {
} }
} }
func TestTeaContext(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
var buf bytes.Buffer
var in bytes.Buffer
m := &testModel{}
p := NewProgram(m, WithContext(ctx), WithInput(&in), WithOutput(&buf))
go func() {
for {
time.Sleep(time.Millisecond)
if m.executed.Load() != nil {
cancel()
return
}
}
}()
if _, err := p.Run(); err != ErrProgramKilled {
t.Fatalf("Expected %v, got %v", ErrProgramKilled, err)
}
}
func TestTeaBatchMsg(t *testing.T) { func TestTeaBatchMsg(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
var in bytes.Buffer var in bytes.Buffer