progress.go 1.54 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
package progress

import (
	"fmt"
	"io"
	"sync"
	"time"
)

type State interface {
	String() string
}

type Progress struct {
15
16
17
	mu sync.Mutex
	w  io.Writer

Michael Yang's avatar
Michael Yang committed
18
19
20
21
22
23
24
	pos int

	ticker *time.Ticker
	states []State
}

func NewProgress(w io.Writer) *Progress {
Michael Yang's avatar
Michael Yang committed
25
	p := &Progress{w: w}
Michael Yang's avatar
Michael Yang committed
26
27
28
29
	go p.start()
	return p
}

Michael Yang's avatar
Michael Yang committed
30
func (p *Progress) Stop() bool {
31
32
33
34
35
36
	for _, state := range p.states {
		if spinner, ok := state.(*Spinner); ok {
			spinner.Stop()
		}
	}

Michael Yang's avatar
Michael Yang committed
37
38
39
40
	if p.ticker != nil {
		p.ticker.Stop()
		p.ticker = nil
		p.render()
41
		fmt.Fprint(p.w, "\n")
Michael Yang's avatar
Michael Yang committed
42
		return true
Michael Yang's avatar
Michael Yang committed
43
	}
Michael Yang's avatar
Michael Yang committed
44
45
46
47
48

	return false
}

func (p *Progress) StopAndClear() bool {
Michael Yang's avatar
Michael Yang committed
49
50
51
	fmt.Fprint(p.w, "\033[?25l")
	defer fmt.Fprint(p.w, "\033[?25h")

Michael Yang's avatar
Michael Yang committed
52
53
54
	stopped := p.Stop()
	if stopped {
		// clear the progress bar by:
Michael Yang's avatar
Michael Yang committed
55
		for i := 0; i < p.pos; i++ {
56
			fmt.Fprint(p.w, "\033[A\033[2K\033[1G")
Michael Yang's avatar
Michael Yang committed
57
		}
Michael Yang's avatar
Michael Yang committed
58
59
60
	}

	return stopped
Michael Yang's avatar
Michael Yang committed
61
62
63
64
65
66
67
68
69
70
71
72
73
}

func (p *Progress) Add(key string, state State) {
	p.mu.Lock()
	defer p.mu.Unlock()

	p.states = append(p.states, state)
}

func (p *Progress) render() error {
	p.mu.Lock()
	defer p.mu.Unlock()

Michael Yang's avatar
Michael Yang committed
74
75
76
	fmt.Fprint(p.w, "\033[?25l")
	defer fmt.Fprint(p.w, "\033[?25h")

77
78
79
80
81
82
	// clear already rendered progress lines
	for i := 0; i < p.pos; i++ {
		if i > 0 {
			fmt.Fprint(p.w, "\033[A")
		}
		fmt.Fprint(p.w, "\033[2K\033[1G")
Michael Yang's avatar
Michael Yang committed
83
84
	}

85
86
87
88
89
90
	// render progress lines
	for i, state := range p.states {
		fmt.Fprint(p.w, state.String())
		if i < len(p.states)-1 {
			fmt.Fprint(p.w, "\n")
		}
Michael Yang's avatar
Michael Yang committed
91
92
	}

93
	p.pos = len(p.states)
Michael Yang's avatar
Michael Yang committed
94
95
96
97
98
99
100
101
102
103

	return nil
}

func (p *Progress) start() {
	p.ticker = time.NewTicker(100 * time.Millisecond)
	for range p.ticker.C {
		p.render()
	}
}