spinner.go 1.35 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
package progress

import (
	"fmt"
	"strings"
6
	"sync/atomic"
Michael Yang's avatar
Michael Yang committed
7
8
9
10
	"time"
)

type Spinner struct {
11
	message      atomic.Value
Michael Yang's avatar
Michael Yang committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
	messageWidth int

	parts []string

	value int

	ticker  *time.Ticker
	started time.Time
	stopped time.Time
}

func NewSpinner(message string) *Spinner {
	s := &Spinner{
		parts: []string{
			"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏",
		},
		started: time.Now(),
	}
30
	s.SetMessage(message)
Michael Yang's avatar
Michael Yang committed
31
32
33
34
	go s.start()
	return s
}

35
36
37
38
func (s *Spinner) SetMessage(message string) {
	s.message.Store(message)
}

Michael Yang's avatar
Michael Yang committed
39
func (s *Spinner) String() string {
40
	var sb strings.Builder
41
42
43

	if message, ok := s.message.Load().(string); ok && len(message) > 0 {
		message := strings.TrimSpace(message)
Michael Yang's avatar
Michael Yang committed
44
45
46
47
		if s.messageWidth > 0 && len(message) > s.messageWidth {
			message = message[:s.messageWidth]
		}

48
		fmt.Fprintf(&sb, "%s", message)
Michael Yang's avatar
Michael Yang committed
49
50
		if padding := s.messageWidth - sb.Len(); padding > 0 {
			sb.WriteString(strings.Repeat(" ", padding))
Michael Yang's avatar
Michael Yang committed
51
52
		}

53
		sb.WriteString(" ")
Michael Yang's avatar
Michael Yang committed
54
55
56
57
	}

	if s.stopped.IsZero() {
		spinner := s.parts[s.value]
58
59
		sb.WriteString(spinner)
		sb.WriteString(" ")
Michael Yang's avatar
Michael Yang committed
60
61
	}

62
	return sb.String()
Michael Yang's avatar
Michael Yang committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
}

func (s *Spinner) start() {
	s.ticker = time.NewTicker(100 * time.Millisecond)
	for range s.ticker.C {
		s.value = (s.value + 1) % len(s.parts)
		if !s.stopped.IsZero() {
			return
		}
	}
}

func (s *Spinner) Stop() {
	if s.stopped.IsZero() {
		s.stopped = time.Now()
	}
}