pb.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588
  1. package pb
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "os"
  7. "strconv"
  8. "strings"
  9. "sync"
  10. "sync/atomic"
  11. "text/template"
  12. "time"
  13. "github.com/fatih/color"
  14. "github.com/mattn/go-colorable"
  15. "github.com/mattn/go-isatty"
  16. "github.com/cheggaaa/pb/v3/termutil"
  17. )
  18. // Version of ProgressBar library
  19. const Version = "3.0.8"
  20. type key int
  21. const (
  22. // Bytes means we're working with byte sizes. Numbers will print as Kb, Mb, etc
  23. // bar.Set(pb.Bytes, true)
  24. Bytes key = 1 << iota
  25. // Use SI bytes prefix names (kB, MB, etc) instead of IEC prefix names (KiB, MiB, etc)
  26. SIBytesPrefix
  27. // Terminal means we're will print to terminal and can use ascii sequences
  28. // Also we're will try to use terminal width
  29. Terminal
  30. // Static means progress bar will not update automaticly
  31. Static
  32. // ReturnSymbol - by default in terminal mode it's '\r'
  33. ReturnSymbol
  34. // Color by default is true when output is tty, but you can set to false for disabling colors
  35. Color
  36. // Hide the progress bar when finished, rather than leaving it up. By default it's false.
  37. CleanOnFinish
  38. // Round elapsed time to this precision. Defaults to time.Second.
  39. TimeRound
  40. )
  41. const (
  42. defaultBarWidth = 100
  43. defaultRefreshRate = time.Millisecond * 200
  44. )
  45. // New creates new ProgressBar object
  46. func New(total int) *ProgressBar {
  47. return New64(int64(total))
  48. }
  49. // New64 creates new ProgressBar object using int64 as total
  50. func New64(total int64) *ProgressBar {
  51. pb := new(ProgressBar)
  52. return pb.SetTotal(total)
  53. }
  54. // StartNew starts new ProgressBar with Default template
  55. func StartNew(total int) *ProgressBar {
  56. return New(total).Start()
  57. }
  58. // Start64 starts new ProgressBar with Default template. Using int64 as total.
  59. func Start64(total int64) *ProgressBar {
  60. return New64(total).Start()
  61. }
  62. var (
  63. terminalWidth = termutil.TerminalWidth
  64. isTerminal = isatty.IsTerminal
  65. isCygwinTerminal = isatty.IsCygwinTerminal
  66. )
  67. // ProgressBar is the main object of bar
  68. type ProgressBar struct {
  69. current, total int64
  70. width int
  71. maxWidth int
  72. mu sync.RWMutex
  73. rm sync.Mutex
  74. vars map[interface{}]interface{}
  75. elements map[string]Element
  76. output io.Writer
  77. coutput io.Writer
  78. nocoutput io.Writer
  79. startTime time.Time
  80. refreshRate time.Duration
  81. tmpl *template.Template
  82. state *State
  83. buf *bytes.Buffer
  84. ticker *time.Ticker
  85. finish chan struct{}
  86. finished bool
  87. configured bool
  88. err error
  89. }
  90. func (pb *ProgressBar) configure() {
  91. if pb.configured {
  92. return
  93. }
  94. pb.configured = true
  95. if pb.vars == nil {
  96. pb.vars = make(map[interface{}]interface{})
  97. }
  98. if pb.output == nil {
  99. pb.output = os.Stderr
  100. }
  101. if pb.tmpl == nil {
  102. pb.tmpl, pb.err = getTemplate(string(Default))
  103. if pb.err != nil {
  104. return
  105. }
  106. }
  107. if pb.vars[Terminal] == nil {
  108. if f, ok := pb.output.(*os.File); ok {
  109. if isTerminal(f.Fd()) || isCygwinTerminal(f.Fd()) {
  110. pb.vars[Terminal] = true
  111. }
  112. }
  113. }
  114. if pb.vars[ReturnSymbol] == nil {
  115. if tm, ok := pb.vars[Terminal].(bool); ok && tm {
  116. pb.vars[ReturnSymbol] = "\r"
  117. }
  118. }
  119. if pb.vars[Color] == nil {
  120. if tm, ok := pb.vars[Terminal].(bool); ok && tm {
  121. pb.vars[Color] = true
  122. }
  123. }
  124. if pb.refreshRate == 0 {
  125. pb.refreshRate = defaultRefreshRate
  126. }
  127. if pb.vars[CleanOnFinish] == nil {
  128. pb.vars[CleanOnFinish] = false
  129. }
  130. if f, ok := pb.output.(*os.File); ok {
  131. pb.coutput = colorable.NewColorable(f)
  132. } else {
  133. pb.coutput = pb.output
  134. }
  135. pb.nocoutput = colorable.NewNonColorable(pb.output)
  136. }
  137. // Start starts the bar
  138. func (pb *ProgressBar) Start() *ProgressBar {
  139. pb.mu.Lock()
  140. defer pb.mu.Unlock()
  141. if pb.finish != nil {
  142. return pb
  143. }
  144. pb.configure()
  145. pb.finished = false
  146. pb.state = nil
  147. pb.startTime = time.Now()
  148. if st, ok := pb.vars[Static].(bool); ok && st {
  149. return pb
  150. }
  151. pb.finish = make(chan struct{})
  152. pb.ticker = time.NewTicker(pb.refreshRate)
  153. go pb.writer(pb.finish)
  154. return pb
  155. }
  156. func (pb *ProgressBar) writer(finish chan struct{}) {
  157. for {
  158. select {
  159. case <-pb.ticker.C:
  160. pb.write(false)
  161. case <-finish:
  162. pb.ticker.Stop()
  163. pb.write(true)
  164. finish <- struct{}{}
  165. return
  166. }
  167. }
  168. }
  169. // Write performs write to the output
  170. func (pb *ProgressBar) Write() *ProgressBar {
  171. pb.mu.RLock()
  172. finished := pb.finished
  173. pb.mu.RUnlock()
  174. pb.write(finished)
  175. return pb
  176. }
  177. func (pb *ProgressBar) write(finish bool) {
  178. result, width := pb.render()
  179. if pb.Err() != nil {
  180. return
  181. }
  182. if pb.GetBool(Terminal) {
  183. if r := (width - CellCount(result)); r > 0 {
  184. result += strings.Repeat(" ", r)
  185. }
  186. }
  187. if ret, ok := pb.Get(ReturnSymbol).(string); ok {
  188. result = ret + result
  189. if finish && ret == "\r" {
  190. if pb.GetBool(CleanOnFinish) {
  191. // "Wipe out" progress bar by overwriting one line with blanks
  192. result = "\r" + color.New(color.Reset).Sprintf(strings.Repeat(" ", width)) + "\r"
  193. } else {
  194. result += "\n"
  195. }
  196. }
  197. }
  198. if pb.GetBool(Color) {
  199. pb.coutput.Write([]byte(result))
  200. } else {
  201. pb.nocoutput.Write([]byte(result))
  202. }
  203. }
  204. // Total return current total bar value
  205. func (pb *ProgressBar) Total() int64 {
  206. return atomic.LoadInt64(&pb.total)
  207. }
  208. // SetTotal sets the total bar value
  209. func (pb *ProgressBar) SetTotal(value int64) *ProgressBar {
  210. atomic.StoreInt64(&pb.total, value)
  211. return pb
  212. }
  213. // AddTotal adds to the total bar value
  214. func (pb *ProgressBar) AddTotal(value int64) *ProgressBar {
  215. atomic.AddInt64(&pb.total, value)
  216. return pb
  217. }
  218. // SetCurrent sets the current bar value
  219. func (pb *ProgressBar) SetCurrent(value int64) *ProgressBar {
  220. atomic.StoreInt64(&pb.current, value)
  221. return pb
  222. }
  223. // Current return current bar value
  224. func (pb *ProgressBar) Current() int64 {
  225. return atomic.LoadInt64(&pb.current)
  226. }
  227. // Add adding given int64 value to bar value
  228. func (pb *ProgressBar) Add64(value int64) *ProgressBar {
  229. atomic.AddInt64(&pb.current, value)
  230. return pb
  231. }
  232. // Add adding given int value to bar value
  233. func (pb *ProgressBar) Add(value int) *ProgressBar {
  234. return pb.Add64(int64(value))
  235. }
  236. // Increment atomically increments the progress
  237. func (pb *ProgressBar) Increment() *ProgressBar {
  238. return pb.Add64(1)
  239. }
  240. // Set sets any value by any key
  241. func (pb *ProgressBar) Set(key, value interface{}) *ProgressBar {
  242. pb.mu.Lock()
  243. defer pb.mu.Unlock()
  244. if pb.vars == nil {
  245. pb.vars = make(map[interface{}]interface{})
  246. }
  247. pb.vars[key] = value
  248. return pb
  249. }
  250. // Get return value by key
  251. func (pb *ProgressBar) Get(key interface{}) interface{} {
  252. pb.mu.RLock()
  253. defer pb.mu.RUnlock()
  254. if pb.vars == nil {
  255. return nil
  256. }
  257. return pb.vars[key]
  258. }
  259. // GetBool return value by key and try to convert there to boolean
  260. // If value doesn't set or not boolean - return false
  261. func (pb *ProgressBar) GetBool(key interface{}) bool {
  262. if v, ok := pb.Get(key).(bool); ok {
  263. return v
  264. }
  265. return false
  266. }
  267. // SetWidth sets the bar width
  268. // When given value <= 0 would be using the terminal width (if possible) or default value.
  269. func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
  270. pb.mu.Lock()
  271. pb.width = width
  272. pb.mu.Unlock()
  273. return pb
  274. }
  275. // SetMaxWidth sets the bar maximum width
  276. // When given value <= 0 would be using the terminal width (if possible) or default value.
  277. func (pb *ProgressBar) SetMaxWidth(maxWidth int) *ProgressBar {
  278. pb.mu.Lock()
  279. pb.maxWidth = maxWidth
  280. pb.mu.Unlock()
  281. return pb
  282. }
  283. // Width return the bar width
  284. // It's current terminal width or settled over 'SetWidth' value.
  285. func (pb *ProgressBar) Width() (width int) {
  286. defer func() {
  287. if r := recover(); r != nil {
  288. width = defaultBarWidth
  289. }
  290. }()
  291. pb.mu.RLock()
  292. width = pb.width
  293. maxWidth := pb.maxWidth
  294. pb.mu.RUnlock()
  295. if width <= 0 {
  296. var err error
  297. if width, err = terminalWidth(); err != nil {
  298. return defaultBarWidth
  299. }
  300. }
  301. if maxWidth > 0 && width > maxWidth {
  302. width = maxWidth
  303. }
  304. return
  305. }
  306. func (pb *ProgressBar) SetRefreshRate(dur time.Duration) *ProgressBar {
  307. pb.mu.Lock()
  308. if dur > 0 {
  309. pb.refreshRate = dur
  310. }
  311. pb.mu.Unlock()
  312. return pb
  313. }
  314. // SetWriter sets the io.Writer. Bar will write in this writer
  315. // By default this is os.Stderr
  316. func (pb *ProgressBar) SetWriter(w io.Writer) *ProgressBar {
  317. pb.mu.Lock()
  318. pb.output = w
  319. pb.configured = false
  320. pb.configure()
  321. pb.mu.Unlock()
  322. return pb
  323. }
  324. // StartTime return the time when bar started
  325. func (pb *ProgressBar) StartTime() time.Time {
  326. pb.mu.RLock()
  327. defer pb.mu.RUnlock()
  328. return pb.startTime
  329. }
  330. // Format convert int64 to string according to the current settings
  331. func (pb *ProgressBar) Format(v int64) string {
  332. if pb.GetBool(Bytes) {
  333. return formatBytes(v, pb.GetBool(SIBytesPrefix))
  334. }
  335. return strconv.FormatInt(v, 10)
  336. }
  337. // Finish stops the bar
  338. func (pb *ProgressBar) Finish() *ProgressBar {
  339. pb.mu.Lock()
  340. if pb.finished {
  341. pb.mu.Unlock()
  342. return pb
  343. }
  344. finishChan := pb.finish
  345. pb.finished = true
  346. pb.mu.Unlock()
  347. if finishChan != nil {
  348. finishChan <- struct{}{}
  349. <-finishChan
  350. pb.mu.Lock()
  351. pb.finish = nil
  352. pb.mu.Unlock()
  353. }
  354. return pb
  355. }
  356. // IsStarted indicates progress bar state
  357. func (pb *ProgressBar) IsStarted() bool {
  358. pb.mu.RLock()
  359. defer pb.mu.RUnlock()
  360. return pb.finish != nil
  361. }
  362. // SetTemplateString sets ProgressBar tempate string and parse it
  363. func (pb *ProgressBar) SetTemplateString(tmpl string) *ProgressBar {
  364. pb.mu.Lock()
  365. defer pb.mu.Unlock()
  366. pb.tmpl, pb.err = getTemplate(tmpl)
  367. return pb
  368. }
  369. // SetTemplateString sets ProgressBarTempate and parse it
  370. func (pb *ProgressBar) SetTemplate(tmpl ProgressBarTemplate) *ProgressBar {
  371. return pb.SetTemplateString(string(tmpl))
  372. }
  373. // NewProxyReader creates a wrapper for given reader, but with progress handle
  374. // Takes io.Reader or io.ReadCloser
  375. // Also, it automatically switches progress bar to handle units as bytes
  376. func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
  377. pb.Set(Bytes, true)
  378. return &Reader{r, pb}
  379. }
  380. // NewProxyWriter creates a wrapper for given writer, but with progress handle
  381. // Takes io.Writer or io.WriteCloser
  382. // Also, it automatically switches progress bar to handle units as bytes
  383. func (pb *ProgressBar) NewProxyWriter(r io.Writer) *Writer {
  384. pb.Set(Bytes, true)
  385. return &Writer{r, pb}
  386. }
  387. func (pb *ProgressBar) render() (result string, width int) {
  388. defer func() {
  389. if r := recover(); r != nil {
  390. pb.SetErr(fmt.Errorf("render panic: %v", r))
  391. }
  392. }()
  393. pb.rm.Lock()
  394. defer pb.rm.Unlock()
  395. pb.mu.Lock()
  396. pb.configure()
  397. if pb.state == nil {
  398. pb.state = &State{ProgressBar: pb}
  399. pb.buf = bytes.NewBuffer(nil)
  400. }
  401. if pb.startTime.IsZero() {
  402. pb.startTime = time.Now()
  403. }
  404. pb.state.id++
  405. pb.state.finished = pb.finished
  406. pb.state.time = time.Now()
  407. pb.mu.Unlock()
  408. pb.state.width = pb.Width()
  409. width = pb.state.width
  410. pb.state.total = pb.Total()
  411. pb.state.current = pb.Current()
  412. pb.buf.Reset()
  413. if e := pb.tmpl.Execute(pb.buf, pb.state); e != nil {
  414. pb.SetErr(e)
  415. return "", 0
  416. }
  417. result = pb.buf.String()
  418. aec := len(pb.state.recalc)
  419. if aec == 0 {
  420. // no adaptive elements
  421. return
  422. }
  423. staticWidth := CellCount(result) - (aec * adElPlaceholderLen)
  424. if pb.state.Width()-staticWidth <= 0 {
  425. result = strings.Replace(result, adElPlaceholder, "", -1)
  426. result = StripString(result, pb.state.Width())
  427. } else {
  428. pb.state.adaptiveElWidth = (width - staticWidth) / aec
  429. for _, el := range pb.state.recalc {
  430. result = strings.Replace(result, adElPlaceholder, el.ProgressElement(pb.state), 1)
  431. }
  432. }
  433. pb.state.recalc = pb.state.recalc[:0]
  434. return
  435. }
  436. // SetErr sets error to the ProgressBar
  437. // Error will be available over Err()
  438. func (pb *ProgressBar) SetErr(err error) *ProgressBar {
  439. pb.mu.Lock()
  440. pb.err = err
  441. pb.mu.Unlock()
  442. return pb
  443. }
  444. // Err return possible error
  445. // When all ok - will be nil
  446. // May contain template.Execute errors
  447. func (pb *ProgressBar) Err() error {
  448. pb.mu.RLock()
  449. defer pb.mu.RUnlock()
  450. return pb.err
  451. }
  452. // String return currrent string representation of ProgressBar
  453. func (pb *ProgressBar) String() string {
  454. res, _ := pb.render()
  455. return res
  456. }
  457. // ProgressElement implements Element interface
  458. func (pb *ProgressBar) ProgressElement(s *State, args ...string) string {
  459. if s.IsAdaptiveWidth() {
  460. pb.SetWidth(s.AdaptiveElWidth())
  461. }
  462. return pb.String()
  463. }
  464. // State represents the current state of bar
  465. // Need for bar elements
  466. type State struct {
  467. *ProgressBar
  468. id uint64
  469. total, current int64
  470. width, adaptiveElWidth int
  471. finished, adaptive bool
  472. time time.Time
  473. recalc []Element
  474. }
  475. // Id it's the current state identifier
  476. // - incremental
  477. // - starts with 1
  478. // - resets after finish/start
  479. func (s *State) Id() uint64 {
  480. return s.id
  481. }
  482. // Total it's bar int64 total
  483. func (s *State) Total() int64 {
  484. return s.total
  485. }
  486. // Value it's current value
  487. func (s *State) Value() int64 {
  488. return s.current
  489. }
  490. // Width of bar
  491. func (s *State) Width() int {
  492. return s.width
  493. }
  494. // AdaptiveElWidth - adaptive elements must return string with given cell count (when AdaptiveElWidth > 0)
  495. func (s *State) AdaptiveElWidth() int {
  496. return s.adaptiveElWidth
  497. }
  498. // IsAdaptiveWidth returns true when element must be shown as adaptive
  499. func (s *State) IsAdaptiveWidth() bool {
  500. return s.adaptive
  501. }
  502. // IsFinished return true when bar is finished
  503. func (s *State) IsFinished() bool {
  504. return s.finished
  505. }
  506. // IsFirst return true only in first render
  507. func (s *State) IsFirst() bool {
  508. return s.id == 1
  509. }
  510. // Time when state was created
  511. func (s *State) Time() time.Time {
  512. return s.time
  513. }