utils/cli/spin/spin.go

package spin

import (
	"context"
	"errors"
	"fmt"

	"github.com/charmbracelet/bubbles/spinner"
	tea "github.com/charmbracelet/bubbletea"
)

type model struct {
	spinner  spinner.Model
	text     string
	quitting bool
	err      error
}

type errMsg error

type textMessage string

func newModel(text string) *model {
	s := spinner.New()
	s.Spinner = spinner.Dot
	return &model{spinner: s, text: text}
}

func (m *model) Init() tea.Cmd {
	return m.spinner.Tick
}

func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case errMsg:
		m.err = msg
		return m, nil

	case textMessage:
		m.text = string(msg)
		return m, nil

	default:
		var cmd tea.Cmd
		m.spinner, cmd = m.spinner.Update(msg)
		return m, cmd
	}
}

func (m *model) View() string {
	if m.err != nil {
		return m.err.Error()
	}
	str := fmt.Sprintf("%s %s\n", m.spinner.View(), m.text)
	if m.quitting {
		return str + "\n"
	}
	return str
}

type Spinner interface {
	Wait() error
	SetMessage(string)
	Err() error
}

var _ Spinner = new(spin)

type spin struct {
	p        *tea.Program
	err      error
	finished chan struct{}
}

func (s *spin) Wait() error {
	<-s.finished
	return s.Err()
}

func (s *spin) SetMessage(msg string) {
	s.p.Send(textMessage(msg))
}

func (s *spin) Err() error {
	switch {
	case errors.Is(s.err, context.Canceled):
		return nil
	case errors.Is(s.err, context.DeadlineExceeded):
		return nil
	case errors.Is(s.err, tea.ErrProgramKilled):
		return nil
	}
	return s.err
}

func Spin(ctx context.Context, message string) Spinner {
	p := tea.NewProgram(
		newModel(message),
		tea.WithContext(ctx),
		tea.WithoutSignalHandler(),
		tea.WithInput(nil),
	)
	s := &spin{p: p, finished: make(chan struct{}, 1)}
	go func() {
		_, err := s.p.Run()
		s.err = err
		s.finished <- struct{}{}
	}()
	return s
}