/
progressbar.go
144 lines (126 loc) · 2.99 KB
/
progressbar.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
package progressbar
import (
"errors"
"fmt"
"io"
"os"
"strings"
"sync"
"time"
)
// ProgressBar is a thread-safe, simple
// progress bar
type ProgressBar struct {
max int // max number of the counter
size int // size of the saucer
currentNum int
currentPercent int
lastPercent int
currentSaucerSize int
lastShown time.Time
startTime time.Time
w io.Writer
theme []string
sync.RWMutex
}
func (p *ProgressBar) SetTheme(theme []string) {
p.Lock()
p.theme = theme
p.Unlock()
}
// New returns a new ProgressBar
// with the specified maximum
func New(max int) *ProgressBar {
return &ProgressBar{
max: max,
size: 40,
theme: []string{"█", " ", "|", "|"},
w: os.Stdout,
lastShown: time.Now(),
startTime: time.Now(),
}
}
// Reset will reset the clock that is used
// to calculate current time and the time left.
func (p *ProgressBar) Reset() {
p.Lock()
defer p.Unlock()
p.lastShown = time.Now()
p.startTime = time.Now()
p.currentNum = 0
}
// SetMax sets the total number of the progress bar
func (p *ProgressBar) SetMax(num int) {
p.Lock()
defer p.Unlock()
p.max = num
}
// SetSize sets the size of the progress bar.
func (p *ProgressBar) SetSize(size int) {
p.Lock()
defer p.Unlock()
p.size = size
}
// SetWriter will specify a different writer than os.Stdout
func (p *ProgressBar) SetWriter(w io.Writer) {
p.Lock()
defer p.Unlock()
p.w = w
}
// Add with increase the current count on the progress bar
func (p *ProgressBar) Add(num int) error {
p.RLock()
currentNum := p.currentNum
p.RUnlock()
return p.Set(currentNum + num)
}
// Set will change the current count on the progress bar
func (p *ProgressBar) Set(num int) error {
p.Lock()
if p.max == 0 {
p.Unlock()
return errors.New("max must be greater than 0")
}
p.currentNum = num
percent := float64(p.currentNum) / float64(p.max)
p.currentSaucerSize = int(percent * float64(p.size))
p.currentPercent = int(percent * 100)
updateBar := p.currentPercent != p.lastPercent && p.currentPercent > 0
p.lastPercent = p.currentPercent
p.Unlock()
if updateBar {
return p.Show()
}
return nil
}
// Show will print the current progress bar
func (p *ProgressBar) Show() error {
p.RLock()
defer p.RUnlock()
if p.currentNum > p.max {
return errors.New("current number exceeds max")
}
s := p.String()
_, err := io.WriteString(p.w, s)
if err != nil {
return err
}
if f, ok := p.w.(*os.File); ok {
f.Sync()
}
return nil
}
func (p *ProgressBar) String() string {
p.RLock()
defer p.RUnlock()
leftTime := time.Since(p.startTime).Seconds() / float64(p.currentNum) * (float64(p.max) - float64(p.currentNum))
return fmt.Sprintf("\r%4d%% %s%s%s%s [%s:%s] ",
p.currentPercent,
p.theme[2],
strings.Repeat(p.theme[0], p.currentSaucerSize),
strings.Repeat(p.theme[1], p.size-p.currentSaucerSize),
p.theme[3],
(time.Duration(time.Since(p.startTime).Seconds()) * time.Second).String(),
(time.Duration(leftTime) * time.Second).String(),
)
}