Skip to content

Commit 5bf66ae

Browse files
astaxieastaxie
committed
support progress bar fix astaxie#14 
1 parent 7f78fbe commit 5bf66ae

File tree

3 files changed

+244
-2
lines changed

3 files changed

+244
-2
lines changed

‎bat.go

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ import (
2828
"os"
2929
"path/filepath"
3030
"runtime"
31+
"strconv"
3132
"strings"
3233
)
3334

@@ -149,15 +150,49 @@ func main() {
149150
}
150151

151152
if download {
152-
_, fl := filepath.Split(u.Path)
153+
var fl string
154+
if disposition := res.Header.Get("Content-Disposition"); disposition != "" {
155+
fls := strings.Split(disposition, ";")
156+
for _, f := range fls {
157+
f = strings.TrimSpace(f)
158+
if strings.HasPrefix(f, "filename=") {
159+
fl = strings.TrimLeft(f, "filename=")
160+
}
161+
}
162+
}
163+
if fl == "" {
164+
_, fl = filepath.Split(u.Path)
165+
}
153166
fd, err := os.OpenFile(fl, os.O_RDWR|os.O_CREATE, 0666)
154167
if err != nil {
155168
log.Fatal("can't create file", err)
156169
}
157-
_, err = io.Copy(fd, res.Body)
170+
if runtime.GOOS != "windowns" {
171+
fmt.Println(Color(res.Proto, Magenta), Color(res.Status, Green))
172+
for k, v := range res.Header {
173+
fmt.Println(Color(k, Gray), ":", Color(strings.Join(v, " "), Cyan))
174+
}
175+
} else {
176+
fmt.Println(res.Proto, res.Status)
177+
for k, v := range res.Header {
178+
fmt.Println(k, ":", strings.Join(v, " "))
179+
}
180+
}
181+
fmt.Println("")
182+
contentLength := res.Header.Get("Content-Length")
183+
var total int64
184+
if contentLength != "" {
185+
total, _ = strconv.ParseInt(contentLength, 10, 64)
186+
}
187+
fmt.Printf("Downloading to \"%s\"\n", fl)
188+
pb := NewProgressBar(total)
189+
pb.Start()
190+
multiWriter := io.MultiWriter(fd, pb)
191+
_, err = io.Copy(multiWriter, res.Body)
158192
if err != nil {
159193
log.Fatal("Can't Write the body into file", err)
160194
}
195+
pb.Finish()
161196
defer fd.Close()
162197
defer res.Body.Close()
163198
return

‎pb.go

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
package main
2+
3+
import (
4+
"fmt"
5+
"math"
6+
"strings"
7+
"sync/atomic"
8+
"time"
9+
)
10+
11+
const (
12+
DEFAULT_REFRESH_RATE = time.Millisecond * 200
13+
FORMAT = "[=>-]"
14+
)
15+
16+
type ProgressBar struct {
17+
current int64
18+
Total int64
19+
RefreshRate time.Duration
20+
ShowPercent, ShowCounters bool
21+
ShowSpeed, ShowTimeLeft, ShowBar bool
22+
ShowFinalTime bool
23+
24+
isFinish int32
25+
startTime time.Time
26+
currentValue int64
27+
28+
BarStart string
29+
BarEnd string
30+
Empty string
31+
Current string
32+
CurrentN string
33+
}
34+
35+
func NewProgressBar(total int64) (pb *ProgressBar) {
36+
pb = &ProgressBar{
37+
Total: total,
38+
RefreshRate: DEFAULT_REFRESH_RATE,
39+
ShowPercent: true,
40+
ShowBar: true,
41+
ShowCounters: true,
42+
ShowFinalTime: true,
43+
ShowTimeLeft: true,
44+
ShowSpeed: true,
45+
BarStart: "[",
46+
BarEnd: "]",
47+
Empty: "_",
48+
Current: "=",
49+
CurrentN: ">",
50+
}
51+
return
52+
}
53+
54+
func (pb *ProgressBar) Start() {
55+
pb.startTime = time.Now()
56+
if pb.Total == 0 {
57+
pb.ShowBar = false
58+
pb.ShowTimeLeft = false
59+
pb.ShowPercent = false
60+
}
61+
go pb.writer()
62+
}
63+
64+
// Write the current state of the progressbar
65+
func (pb *ProgressBar) Update() {
66+
c := atomic.LoadInt64(&pb.current)
67+
if c != pb.currentValue {
68+
pb.write(c)
69+
pb.currentValue = c
70+
}
71+
}
72+
73+
// Internal loop for writing progressbar
74+
func (pb *ProgressBar) writer() {
75+
for {
76+
if atomic.LoadInt32(&pb.isFinish) != 0 {
77+
break
78+
}
79+
pb.Update()
80+
time.Sleep(pb.RefreshRate)
81+
}
82+
}
83+
84+
// Increment current value
85+
func (pb *ProgressBar) Increment() int {
86+
return pb.Add(1)
87+
}
88+
89+
// Set current value
90+
func (pb *ProgressBar) Set(current int) {
91+
atomic.StoreInt64(&pb.current, int64(current))
92+
}
93+
94+
// Add to current value
95+
func (pb *ProgressBar) Add(add int) int {
96+
return int(pb.Add64(int64(add)))
97+
}
98+
99+
func (pb *ProgressBar) Add64(add int64) int64 {
100+
return atomic.AddInt64(&pb.current, add)
101+
}
102+
103+
// End print
104+
func (pb *ProgressBar) Finish() {
105+
atomic.StoreInt32(&pb.isFinish, 1)
106+
pb.write(atomic.LoadInt64(&pb.current))
107+
}
108+
109+
// implement io.Writer
110+
func (pb *ProgressBar) Write(p []byte) (n int, err error) {
111+
n = len(p)
112+
pb.Add(n)
113+
return
114+
}
115+
116+
func (pb *ProgressBar) write(current int64) {
117+
width := 123
118+
119+
var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string
120+
121+
// percents
122+
if pb.ShowPercent {
123+
percent := float64(current) / (float64(pb.Total) / float64(100))
124+
percentBox = fmt.Sprintf(" %#.02f %% ", percent)
125+
}
126+
127+
// counters
128+
if pb.ShowCounters {
129+
if pb.Total > 0 {
130+
countersBox = fmt.Sprintf("%s / %s ", FormatBytes(current), FormatBytes(pb.Total))
131+
} else {
132+
countersBox = FormatBytes(current) + " "
133+
}
134+
}
135+
136+
// time left
137+
fromStart := time.Now().Sub(pb.startTime)
138+
if atomic.LoadInt32(&pb.isFinish) != 0 {
139+
if pb.ShowFinalTime {
140+
left := (fromStart / time.Second) * time.Second
141+
timeLeftBox = left.String()
142+
}
143+
} else if pb.ShowTimeLeft && current > 0 {
144+
perEntry := fromStart / time.Duration(current)
145+
left := time.Duration(pb.Total-current) * perEntry
146+
left = (left / time.Second) * time.Second
147+
timeLeftBox = left.String()
148+
}
149+
150+
// speed
151+
if pb.ShowSpeed && current > 0 {
152+
fromStart := time.Now().Sub(pb.startTime)
153+
speed := float64(current) / (float64(fromStart) / float64(time.Second))
154+
speedBox = FormatBytes(int64(speed)) + "/s "
155+
}
156+
157+
// bar
158+
if pb.ShowBar {
159+
size := width - len(countersBox+pb.BarStart+pb.BarEnd+percentBox+timeLeftBox+speedBox)
160+
if size > 0 {
161+
curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
162+
emptCount := size - curCount
163+
barBox = pb.BarStart
164+
if emptCount < 0 {
165+
emptCount = 0
166+
}
167+
if curCount > size {
168+
curCount = size
169+
}
170+
if emptCount <= 0 {
171+
barBox += strings.Repeat(pb.Current, curCount)
172+
} else if curCount > 0 {
173+
barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
174+
}
175+
176+
barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
177+
}
178+
}
179+
180+
// check len
181+
out = countersBox + barBox + percentBox + speedBox + timeLeftBox
182+
if len(out) < width {
183+
end = strings.Repeat(" ", width-len(out))
184+
}
185+
186+
// and print!
187+
fmt.Print("\r" + out + end + "\n")
188+
}

‎utils.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"fmt"
45
"strconv"
56
"strings"
67
)
@@ -59,3 +60,21 @@ func isbool(v string) (bool, error) {
5960
func isfloat(v string) (float64, error) {
6061
return strconv.ParseFloat(v, 64)
6162
}
63+
64+
// Convert bytes to human readable string. Like a 2 MB, 64.2 KB, 52 B
65+
func FormatBytes(i int64) (result string) {
66+
switch {
67+
case i > (1024 * 1024 * 1024 * 1024):
68+
result = fmt.Sprintf("%#.02f TB", float64(i)/1024/1024/1024/1024)
69+
case i > (1024 * 1024 * 1024):
70+
result = fmt.Sprintf("%#.02f GB", float64(i)/1024/1024/1024)
71+
case i > (1024 * 1024):
72+
result = fmt.Sprintf("%#.02f MB", float64(i)/1024/1024)
73+
case i > 1024:
74+
result = fmt.Sprintf("%#.02f KB", float64(i)/1024)
75+
default:
76+
result = fmt.Sprintf("%d B", i)
77+
}
78+
result = strings.Trim(result, " ")
79+
return
80+
}

0 commit comments

Comments
 (0)