pb.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500
  1. // Simple console progress bars
  2. package pb
  3. import (
  4. "fmt"
  5. "io"
  6. "math"
  7. "strings"
  8. "sync"
  9. "sync/atomic"
  10. "time"
  11. "unicode/utf8"
  12. )
  13. // Current version
  14. const Version = "1.0.25"
  15. const (
  16. // Default refresh rate - 200ms
  17. DEFAULT_REFRESH_RATE = time.Millisecond * 200
  18. FORMAT = "[=>-]"
  19. )
  20. // DEPRECATED
  21. // variables for backward compatibility, from now do not work
  22. // use pb.Format and pb.SetRefreshRate
  23. var (
  24. DefaultRefreshRate = DEFAULT_REFRESH_RATE
  25. BarStart, BarEnd, Empty, Current, CurrentN string
  26. )
  27. // Create new progress bar object
  28. func New(total int) *ProgressBar {
  29. return New64(int64(total))
  30. }
  31. // Create new progress bar object using int64 as total
  32. func New64(total int64) *ProgressBar {
  33. pb := &ProgressBar{
  34. Total: total,
  35. RefreshRate: DEFAULT_REFRESH_RATE,
  36. ShowPercent: true,
  37. ShowCounters: true,
  38. ShowBar: true,
  39. ShowTimeLeft: true,
  40. ShowElapsedTime: false,
  41. ShowFinalTime: true,
  42. Units: U_NO,
  43. ManualUpdate: false,
  44. finish: make(chan struct{}),
  45. }
  46. return pb.Format(FORMAT)
  47. }
  48. // Create new object and start
  49. func StartNew(total int) *ProgressBar {
  50. return New(total).Start()
  51. }
  52. // Callback for custom output
  53. // For example:
  54. // bar.Callback = func(s string) {
  55. // mySuperPrint(s)
  56. // }
  57. //
  58. type Callback func(out string)
  59. type ProgressBar struct {
  60. current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
  61. previous int64
  62. Total int64
  63. RefreshRate time.Duration
  64. ShowPercent, ShowCounters bool
  65. ShowSpeed, ShowTimeLeft, ShowBar bool
  66. ShowFinalTime, ShowElapsedTime bool
  67. Output io.Writer
  68. Callback Callback
  69. NotPrint bool
  70. Units Units
  71. Width int
  72. ForceWidth bool
  73. ManualUpdate bool
  74. AutoStat bool
  75. // Default width for the time box.
  76. UnitsWidth int
  77. TimeBoxWidth int
  78. finishOnce sync.Once //Guards isFinish
  79. finish chan struct{}
  80. isFinish bool
  81. startTime time.Time
  82. startValue int64
  83. changeTime time.Time
  84. prefix, postfix string
  85. mu sync.Mutex
  86. lastPrint string
  87. BarStart string
  88. BarEnd string
  89. Empty string
  90. Current string
  91. CurrentN string
  92. AlwaysUpdate bool
  93. }
  94. // Start print
  95. func (pb *ProgressBar) Start() *ProgressBar {
  96. pb.startTime = time.Now()
  97. pb.startValue = atomic.LoadInt64(&pb.current)
  98. if atomic.LoadInt64(&pb.Total) == 0 {
  99. pb.ShowTimeLeft = false
  100. pb.ShowPercent = false
  101. pb.AutoStat = false
  102. }
  103. if !pb.ManualUpdate {
  104. pb.Update() // Initial printing of the bar before running the bar refresher.
  105. go pb.refresher()
  106. }
  107. return pb
  108. }
  109. // Increment current value
  110. func (pb *ProgressBar) Increment() int {
  111. return pb.Add(1)
  112. }
  113. // Get current value
  114. func (pb *ProgressBar) Get() int64 {
  115. c := atomic.LoadInt64(&pb.current)
  116. return c
  117. }
  118. // Set current value
  119. func (pb *ProgressBar) Set(current int) *ProgressBar {
  120. return pb.Set64(int64(current))
  121. }
  122. // Set64 sets the current value as int64
  123. func (pb *ProgressBar) Set64(current int64) *ProgressBar {
  124. atomic.StoreInt64(&pb.current, current)
  125. return pb
  126. }
  127. // Add to current value
  128. func (pb *ProgressBar) Add(add int) int {
  129. return int(pb.Add64(int64(add)))
  130. }
  131. func (pb *ProgressBar) Add64(add int64) int64 {
  132. return atomic.AddInt64(&pb.current, add)
  133. }
  134. // Set prefix string
  135. func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
  136. pb.prefix = prefix
  137. return pb
  138. }
  139. // Set postfix string
  140. func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
  141. pb.postfix = postfix
  142. return pb
  143. }
  144. // Set custom format for bar
  145. // Example: bar.Format("[=>_]")
  146. // Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
  147. func (pb *ProgressBar) Format(format string) *ProgressBar {
  148. var formatEntries []string
  149. if utf8.RuneCountInString(format) == 5 {
  150. formatEntries = strings.Split(format, "")
  151. } else {
  152. formatEntries = strings.Split(format, "\x00")
  153. }
  154. if len(formatEntries) == 5 {
  155. pb.BarStart = formatEntries[0]
  156. pb.BarEnd = formatEntries[4]
  157. pb.Empty = formatEntries[3]
  158. pb.Current = formatEntries[1]
  159. pb.CurrentN = formatEntries[2]
  160. }
  161. return pb
  162. }
  163. // Set bar refresh rate
  164. func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
  165. pb.RefreshRate = rate
  166. return pb
  167. }
  168. // Set units
  169. // bar.SetUnits(U_NO) - by default
  170. // bar.SetUnits(U_BYTES) - for Mb, Kb, etc
  171. func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
  172. pb.Units = units
  173. return pb
  174. }
  175. // Set max width, if width is bigger than terminal width, will be ignored
  176. func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
  177. pb.Width = width
  178. pb.ForceWidth = false
  179. return pb
  180. }
  181. // Set bar width
  182. func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
  183. pb.Width = width
  184. pb.ForceWidth = true
  185. return pb
  186. }
  187. // End print
  188. func (pb *ProgressBar) Finish() {
  189. //Protect multiple calls
  190. pb.finishOnce.Do(func() {
  191. close(pb.finish)
  192. pb.write(atomic.LoadInt64(&pb.Total), atomic.LoadInt64(&pb.current))
  193. pb.mu.Lock()
  194. defer pb.mu.Unlock()
  195. switch {
  196. case pb.Output != nil:
  197. fmt.Fprintln(pb.Output)
  198. case !pb.NotPrint:
  199. fmt.Println()
  200. }
  201. pb.isFinish = true
  202. })
  203. }
  204. // IsFinished return boolean
  205. func (pb *ProgressBar) IsFinished() bool {
  206. pb.mu.Lock()
  207. defer pb.mu.Unlock()
  208. return pb.isFinish
  209. }
  210. // End print and write string 'str'
  211. func (pb *ProgressBar) FinishPrint(str string) {
  212. pb.Finish()
  213. if pb.Output != nil {
  214. fmt.Fprintln(pb.Output, str)
  215. } else {
  216. fmt.Println(str)
  217. }
  218. }
  219. // implement io.Writer
  220. func (pb *ProgressBar) Write(p []byte) (n int, err error) {
  221. n = len(p)
  222. pb.Add(n)
  223. return
  224. }
  225. // implement io.Reader
  226. func (pb *ProgressBar) Read(p []byte) (n int, err error) {
  227. n = len(p)
  228. pb.Add(n)
  229. return
  230. }
  231. // Create new proxy reader over bar
  232. // Takes io.Reader or io.ReadCloser
  233. func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
  234. return &Reader{r, pb}
  235. }
  236. func (pb *ProgressBar) write(total, current int64) {
  237. width := pb.GetWidth()
  238. var percentBox, countersBox, timeLeftBox, timeSpentBox, speedBox, barBox, end, out string
  239. // percents
  240. if pb.ShowPercent {
  241. var percent float64
  242. if total > 0 {
  243. percent = float64(current) / (float64(total) / float64(100))
  244. } else {
  245. percent = float64(current) / float64(100)
  246. }
  247. percentBox = fmt.Sprintf(" %6.02f%%", percent)
  248. }
  249. // counters
  250. if pb.ShowCounters {
  251. current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
  252. if total > 0 {
  253. totalS := Format(total).To(pb.Units).Width(pb.UnitsWidth)
  254. countersBox = fmt.Sprintf(" %s / %s ", current, totalS)
  255. } else {
  256. countersBox = fmt.Sprintf(" %s / ? ", current)
  257. }
  258. }
  259. // time left
  260. pb.mu.Lock()
  261. currentFromStart := current - pb.startValue
  262. fromStart := time.Now().Sub(pb.startTime)
  263. lastChangeTime := pb.changeTime
  264. fromChange := lastChangeTime.Sub(pb.startTime)
  265. pb.mu.Unlock()
  266. if pb.ShowElapsedTime {
  267. timeSpentBox = fmt.Sprintf(" %s ", (fromStart/time.Second)*time.Second)
  268. }
  269. select {
  270. case <-pb.finish:
  271. if pb.ShowFinalTime {
  272. var left time.Duration
  273. left = (fromStart / time.Second) * time.Second
  274. timeLeftBox = fmt.Sprintf(" %s", left.String())
  275. }
  276. default:
  277. if pb.ShowTimeLeft && currentFromStart > 0 {
  278. perEntry := fromChange / time.Duration(currentFromStart)
  279. var left time.Duration
  280. if total > 0 {
  281. left = time.Duration(total-currentFromStart) * perEntry
  282. left -= time.Since(lastChangeTime)
  283. left = (left / time.Second) * time.Second
  284. } else {
  285. left = time.Duration(currentFromStart) * perEntry
  286. left = (left / time.Second) * time.Second
  287. }
  288. if left > 0 {
  289. timeLeft := Format(int64(left)).To(U_DURATION).String()
  290. timeLeftBox = fmt.Sprintf(" %s", timeLeft)
  291. }
  292. }
  293. }
  294. if len(timeLeftBox) < pb.TimeBoxWidth {
  295. timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
  296. }
  297. // speed
  298. if pb.ShowSpeed && currentFromStart > 0 {
  299. fromStart := time.Now().Sub(pb.startTime)
  300. speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
  301. speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
  302. }
  303. barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeSpentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
  304. // bar
  305. if pb.ShowBar {
  306. size := width - barWidth
  307. if size > 0 {
  308. if total > 0 {
  309. curSize := int(math.Ceil((float64(current) / float64(total)) * float64(size)))
  310. emptySize := size - curSize
  311. barBox = pb.BarStart
  312. if emptySize < 0 {
  313. emptySize = 0
  314. }
  315. if curSize > size {
  316. curSize = size
  317. }
  318. cursorLen := escapeAwareRuneCountInString(pb.Current)
  319. if emptySize <= 0 {
  320. barBox += strings.Repeat(pb.Current, curSize/cursorLen)
  321. } else if curSize > 0 {
  322. cursorEndLen := escapeAwareRuneCountInString(pb.CurrentN)
  323. cursorRepetitions := (curSize - cursorEndLen) / cursorLen
  324. barBox += strings.Repeat(pb.Current, cursorRepetitions)
  325. barBox += pb.CurrentN
  326. }
  327. emptyLen := escapeAwareRuneCountInString(pb.Empty)
  328. barBox += strings.Repeat(pb.Empty, emptySize/emptyLen)
  329. barBox += pb.BarEnd
  330. } else {
  331. pos := size - int(current)%int(size)
  332. barBox = pb.BarStart
  333. if pos-1 > 0 {
  334. barBox += strings.Repeat(pb.Empty, pos-1)
  335. }
  336. barBox += pb.Current
  337. if size-pos-1 > 0 {
  338. barBox += strings.Repeat(pb.Empty, size-pos-1)
  339. }
  340. barBox += pb.BarEnd
  341. }
  342. }
  343. }
  344. // check len
  345. out = pb.prefix + timeSpentBox + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
  346. if cl := escapeAwareRuneCountInString(out); cl < width {
  347. end = strings.Repeat(" ", width-cl)
  348. }
  349. // and print!
  350. pb.mu.Lock()
  351. defer pb.mu.Unlock()
  352. pb.lastPrint = out + end
  353. isFinish := pb.isFinish
  354. switch {
  355. case isFinish:
  356. return
  357. case pb.Output != nil:
  358. fmt.Fprint(pb.Output, "\r"+out+end)
  359. case pb.Callback != nil:
  360. pb.Callback(out + end)
  361. case !pb.NotPrint:
  362. fmt.Print("\r" + out + end)
  363. }
  364. }
  365. // GetTerminalWidth - returns terminal width for all platforms.
  366. func GetTerminalWidth() (int, error) {
  367. return terminalWidth()
  368. }
  369. func (pb *ProgressBar) GetWidth() int {
  370. if pb.ForceWidth {
  371. return pb.Width
  372. }
  373. width := pb.Width
  374. termWidth, _ := terminalWidth()
  375. if width == 0 || termWidth <= width {
  376. width = termWidth
  377. }
  378. return width
  379. }
  380. // Write the current state of the progressbar
  381. func (pb *ProgressBar) Update() {
  382. c := atomic.LoadInt64(&pb.current)
  383. p := atomic.LoadInt64(&pb.previous)
  384. t := atomic.LoadInt64(&pb.Total)
  385. if p != c {
  386. pb.mu.Lock()
  387. pb.changeTime = time.Now()
  388. pb.mu.Unlock()
  389. atomic.StoreInt64(&pb.previous, c)
  390. }
  391. pb.write(t, c)
  392. if pb.AutoStat {
  393. if c == 0 {
  394. pb.startTime = time.Now()
  395. pb.startValue = 0
  396. } else if c >= t && pb.isFinish != true {
  397. pb.Finish()
  398. }
  399. }
  400. }
  401. // String return the last bar print
  402. func (pb *ProgressBar) String() string {
  403. pb.mu.Lock()
  404. defer pb.mu.Unlock()
  405. return pb.lastPrint
  406. }
  407. // SetTotal atomically sets new total count
  408. func (pb *ProgressBar) SetTotal(total int) *ProgressBar {
  409. return pb.SetTotal64(int64(total))
  410. }
  411. // SetTotal64 atomically sets new total count
  412. func (pb *ProgressBar) SetTotal64(total int64) *ProgressBar {
  413. atomic.StoreInt64(&pb.Total, total)
  414. return pb
  415. }
  416. // Reset bar and set new total count
  417. // Does effect only on finished bar
  418. func (pb *ProgressBar) Reset(total int) *ProgressBar {
  419. pb.mu.Lock()
  420. defer pb.mu.Unlock()
  421. if pb.isFinish {
  422. pb.SetTotal(total).Set(0)
  423. atomic.StoreInt64(&pb.previous, 0)
  424. }
  425. return pb
  426. }
  427. // Internal loop for refreshing the progressbar
  428. func (pb *ProgressBar) refresher() {
  429. for {
  430. select {
  431. case <-pb.finish:
  432. return
  433. case <-time.After(pb.RefreshRate):
  434. pb.Update()
  435. }
  436. }
  437. }