pool.go 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. // +build linux darwin freebsd netbsd openbsd solaris dragonfly
  2. package pb
  3. import (
  4. "fmt"
  5. "sync"
  6. "time"
  7. )
  8. // Create and start new pool with given bars
  9. // You need call pool.Stop() after work
  10. func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
  11. pool = new(Pool)
  12. if err = pool.start(); err != nil {
  13. return
  14. }
  15. pool.add(pbs...)
  16. return
  17. }
  18. type Pool struct {
  19. RefreshRate time.Duration
  20. bars []*ProgressBar
  21. quit chan int
  22. finishOnce sync.Once
  23. }
  24. func (p *Pool) add(pbs ...*ProgressBar) {
  25. for _, bar := range pbs {
  26. bar.ManualUpdate = true
  27. bar.NotPrint = true
  28. bar.Start()
  29. p.bars = append(p.bars, bar)
  30. }
  31. }
  32. func (p *Pool) start() (err error) {
  33. p.RefreshRate = DefaultRefreshRate
  34. quit, err := lockEcho()
  35. if err != nil {
  36. return
  37. }
  38. p.quit = make(chan int)
  39. go p.writer(quit)
  40. return
  41. }
  42. func (p *Pool) writer(finish chan int) {
  43. var first = true
  44. for {
  45. select {
  46. case <-time.After(p.RefreshRate):
  47. if p.print(first) {
  48. p.print(false)
  49. finish <- 1
  50. return
  51. }
  52. first = false
  53. case <-p.quit:
  54. finish <- 1
  55. return
  56. }
  57. }
  58. }
  59. func (p *Pool) print(first bool) bool {
  60. var out string
  61. if !first {
  62. out = fmt.Sprintf("\033[%dA", len(p.bars))
  63. }
  64. isFinished := true
  65. for _, bar := range p.bars {
  66. if !bar.isFinish {
  67. isFinished = false
  68. }
  69. bar.Update()
  70. out += fmt.Sprintf("\r%s\n", bar.String())
  71. }
  72. fmt.Print(out)
  73. return isFinished
  74. }
  75. // Restore terminal state and close pool
  76. func (p *Pool) Stop() error {
  77. // Wait until one final refresh has passed.
  78. time.Sleep(p.RefreshRate)
  79. p.finishOnce.Do(func() {
  80. close(p.quit)
  81. })
  82. return unlockEcho()
  83. }