pb.go 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424
  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.2"
  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. ShowFinalTime: true,
  41. Units: U_NO,
  42. ManualUpdate: false,
  43. finish: make(chan struct{}),
  44. currentValue: -1,
  45. mu: new(sync.Mutex),
  46. }
  47. return pb.Format(FORMAT)
  48. }
  49. // Create new object and start
  50. func StartNew(total int) *ProgressBar {
  51. return New(total).Start()
  52. }
  53. // Callback for custom output
  54. // For example:
  55. // bar.Callback = func(s string) {
  56. // mySuperPrint(s)
  57. // }
  58. //
  59. type Callback func(out string)
  60. type ProgressBar struct {
  61. current int64 // current must be first member of struct (https://code.google.com/p/go/issues/detail?id=5278)
  62. Total int64
  63. RefreshRate time.Duration
  64. ShowPercent, ShowCounters bool
  65. ShowSpeed, ShowTimeLeft, ShowBar bool
  66. ShowFinalTime bool
  67. Output io.Writer
  68. Callback Callback
  69. NotPrint bool
  70. Units Units
  71. Width int
  72. ForceWidth bool
  73. ManualUpdate bool
  74. // Default width for the time box.
  75. UnitsWidth int
  76. TimeBoxWidth int
  77. finishOnce sync.Once //Guards isFinish
  78. finish chan struct{}
  79. isFinish bool
  80. startTime time.Time
  81. startValue int64
  82. currentValue int64
  83. prefix, postfix string
  84. mu *sync.Mutex
  85. lastPrint string
  86. BarStart string
  87. BarEnd string
  88. Empty string
  89. Current string
  90. CurrentN string
  91. AlwaysUpdate bool
  92. }
  93. // Start print
  94. func (pb *ProgressBar) Start() *ProgressBar {
  95. pb.startTime = time.Now()
  96. pb.startValue = pb.current
  97. if pb.Total == 0 {
  98. pb.ShowTimeLeft = false
  99. pb.ShowPercent = false
  100. }
  101. if !pb.ManualUpdate {
  102. pb.Update() // Initial printing of the bar before running the bar refresher.
  103. go pb.refresher()
  104. }
  105. return pb
  106. }
  107. // Increment current value
  108. func (pb *ProgressBar) Increment() int {
  109. return pb.Add(1)
  110. }
  111. // Set current value
  112. func (pb *ProgressBar) Set(current int) *ProgressBar {
  113. return pb.Set64(int64(current))
  114. }
  115. // Set64 sets the current value as int64
  116. func (pb *ProgressBar) Set64(current int64) *ProgressBar {
  117. atomic.StoreInt64(&pb.current, current)
  118. return pb
  119. }
  120. // Add to current value
  121. func (pb *ProgressBar) Add(add int) int {
  122. return int(pb.Add64(int64(add)))
  123. }
  124. func (pb *ProgressBar) Add64(add int64) int64 {
  125. return atomic.AddInt64(&pb.current, add)
  126. }
  127. // Set prefix string
  128. func (pb *ProgressBar) Prefix(prefix string) *ProgressBar {
  129. pb.prefix = prefix
  130. return pb
  131. }
  132. // Set postfix string
  133. func (pb *ProgressBar) Postfix(postfix string) *ProgressBar {
  134. pb.postfix = postfix
  135. return pb
  136. }
  137. // Set custom format for bar
  138. // Example: bar.Format("[=>_]")
  139. // Example: bar.Format("[\x00=\x00>\x00-\x00]") // \x00 is the delimiter
  140. func (pb *ProgressBar) Format(format string) *ProgressBar {
  141. var formatEntries []string
  142. if len(format) == 5 {
  143. formatEntries = strings.Split(format, "")
  144. } else {
  145. formatEntries = strings.Split(format, "\x00")
  146. }
  147. if len(formatEntries) == 5 {
  148. pb.BarStart = formatEntries[0]
  149. pb.BarEnd = formatEntries[4]
  150. pb.Empty = formatEntries[3]
  151. pb.Current = formatEntries[1]
  152. pb.CurrentN = formatEntries[2]
  153. }
  154. return pb
  155. }
  156. // Set bar refresh rate
  157. func (pb *ProgressBar) SetRefreshRate(rate time.Duration) *ProgressBar {
  158. pb.RefreshRate = rate
  159. return pb
  160. }
  161. // Set units
  162. // bar.SetUnits(U_NO) - by default
  163. // bar.SetUnits(U_BYTES) - for Mb, Kb, etc
  164. func (pb *ProgressBar) SetUnits(units Units) *ProgressBar {
  165. pb.Units = units
  166. return pb
  167. }
  168. // Set max width, if width is bigger than terminal width, will be ignored
  169. func (pb *ProgressBar) SetMaxWidth(width int) *ProgressBar {
  170. pb.Width = width
  171. pb.ForceWidth = false
  172. return pb
  173. }
  174. // Set bar width
  175. func (pb *ProgressBar) SetWidth(width int) *ProgressBar {
  176. pb.Width = width
  177. pb.ForceWidth = true
  178. return pb
  179. }
  180. // End print
  181. func (pb *ProgressBar) Finish() {
  182. //Protect multiple calls
  183. pb.finishOnce.Do(func() {
  184. close(pb.finish)
  185. pb.write(atomic.LoadInt64(&pb.current))
  186. if !pb.NotPrint {
  187. fmt.Println()
  188. }
  189. pb.isFinish = true
  190. })
  191. }
  192. // End print and write string 'str'
  193. func (pb *ProgressBar) FinishPrint(str string) {
  194. pb.Finish()
  195. fmt.Println(str)
  196. }
  197. // implement io.Writer
  198. func (pb *ProgressBar) Write(p []byte) (n int, err error) {
  199. n = len(p)
  200. pb.Add(n)
  201. return
  202. }
  203. // implement io.Reader
  204. func (pb *ProgressBar) Read(p []byte) (n int, err error) {
  205. n = len(p)
  206. pb.Add(n)
  207. return
  208. }
  209. // Create new proxy reader over bar
  210. func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader {
  211. return &Reader{r, pb}
  212. }
  213. func (pb *ProgressBar) write(current int64) {
  214. width := pb.GetWidth()
  215. var percentBox, countersBox, timeLeftBox, speedBox, barBox, end, out string
  216. // percents
  217. if pb.ShowPercent {
  218. var percent float64
  219. if pb.Total > 0 {
  220. percent = float64(current) / (float64(pb.Total) / float64(100))
  221. } else {
  222. percent = float64(current) / float64(100)
  223. }
  224. percentBox = fmt.Sprintf(" %6.02f%%", percent)
  225. }
  226. // counters
  227. if pb.ShowCounters {
  228. current := Format(current).To(pb.Units).Width(pb.UnitsWidth)
  229. if pb.Total > 0 {
  230. total := Format(pb.Total).To(pb.Units).Width(pb.UnitsWidth)
  231. countersBox = fmt.Sprintf(" %s / %s ", current, total)
  232. } else {
  233. countersBox = fmt.Sprintf(" %s / ? ", current)
  234. }
  235. }
  236. // time left
  237. fromStart := time.Now().Sub(pb.startTime)
  238. currentFromStart := current - pb.startValue
  239. select {
  240. case <-pb.finish:
  241. if pb.ShowFinalTime {
  242. var left time.Duration
  243. if pb.Total > 0 {
  244. left = (fromStart / time.Second) * time.Second
  245. } else {
  246. left = (time.Duration(currentFromStart) / time.Second) * time.Second
  247. }
  248. timeLeftBox = left.String()
  249. }
  250. default:
  251. if pb.ShowTimeLeft && currentFromStart > 0 {
  252. perEntry := fromStart / time.Duration(currentFromStart)
  253. var left time.Duration
  254. if pb.Total > 0 {
  255. left = time.Duration(pb.Total-currentFromStart) * perEntry
  256. left = (left / time.Second) * time.Second
  257. } else {
  258. left = time.Duration(currentFromStart) * perEntry
  259. left = (left / time.Second) * time.Second
  260. }
  261. timeLeft := Format(int64(left)).To(U_DURATION).String()
  262. timeLeftBox = fmt.Sprintf(" %s", timeLeft)
  263. }
  264. }
  265. if len(timeLeftBox) < pb.TimeBoxWidth {
  266. timeLeftBox = fmt.Sprintf("%s%s", strings.Repeat(" ", pb.TimeBoxWidth-len(timeLeftBox)), timeLeftBox)
  267. }
  268. // speed
  269. if pb.ShowSpeed && currentFromStart > 0 {
  270. fromStart := time.Now().Sub(pb.startTime)
  271. speed := float64(currentFromStart) / (float64(fromStart) / float64(time.Second))
  272. speedBox = " " + Format(int64(speed)).To(pb.Units).Width(pb.UnitsWidth).PerSec().String()
  273. }
  274. barWidth := escapeAwareRuneCountInString(countersBox + pb.BarStart + pb.BarEnd + percentBox + timeLeftBox + speedBox + pb.prefix + pb.postfix)
  275. // bar
  276. if pb.ShowBar {
  277. size := width - barWidth
  278. if size > 0 {
  279. if pb.Total > 0 {
  280. curCount := int(math.Ceil((float64(current) / float64(pb.Total)) * float64(size)))
  281. emptCount := size - curCount
  282. barBox = pb.BarStart
  283. if emptCount < 0 {
  284. emptCount = 0
  285. }
  286. if curCount > size {
  287. curCount = size
  288. }
  289. if emptCount <= 0 {
  290. barBox += strings.Repeat(pb.Current, curCount)
  291. } else if curCount > 0 {
  292. barBox += strings.Repeat(pb.Current, curCount-1) + pb.CurrentN
  293. }
  294. barBox += strings.Repeat(pb.Empty, emptCount) + pb.BarEnd
  295. } else {
  296. barBox = pb.BarStart
  297. pos := size - int(current)%int(size)
  298. if pos-1 > 0 {
  299. barBox += strings.Repeat(pb.Empty, pos-1)
  300. }
  301. barBox += pb.Current
  302. if size-pos-1 > 0 {
  303. barBox += strings.Repeat(pb.Empty, size-pos-1)
  304. }
  305. barBox += pb.BarEnd
  306. }
  307. }
  308. }
  309. // check len
  310. out = pb.prefix + countersBox + barBox + percentBox + speedBox + timeLeftBox + pb.postfix
  311. if escapeAwareRuneCountInString(out) < width {
  312. end = strings.Repeat(" ", width-utf8.RuneCountInString(out))
  313. }
  314. // and print!
  315. pb.mu.Lock()
  316. pb.lastPrint = out + end
  317. pb.mu.Unlock()
  318. switch {
  319. case pb.isFinish:
  320. return
  321. case pb.Output != nil:
  322. fmt.Fprint(pb.Output, "\r"+out+end)
  323. case pb.Callback != nil:
  324. pb.Callback(out + end)
  325. case !pb.NotPrint:
  326. fmt.Print("\r" + out + end)
  327. }
  328. }
  329. // GetTerminalWidth - returns terminal width for all platforms.
  330. func GetTerminalWidth() (int, error) {
  331. return terminalWidth()
  332. }
  333. func (pb *ProgressBar) GetWidth() int {
  334. if pb.ForceWidth {
  335. return pb.Width
  336. }
  337. width := pb.Width
  338. termWidth, _ := terminalWidth()
  339. if width == 0 || termWidth <= width {
  340. width = termWidth
  341. }
  342. return width
  343. }
  344. // Write the current state of the progressbar
  345. func (pb *ProgressBar) Update() {
  346. c := atomic.LoadInt64(&pb.current)
  347. if pb.AlwaysUpdate || c != pb.currentValue {
  348. pb.write(c)
  349. pb.currentValue = c
  350. }
  351. }
  352. func (pb *ProgressBar) String() string {
  353. return pb.lastPrint
  354. }
  355. // Internal loop for refreshing the progressbar
  356. func (pb *ProgressBar) refresher() {
  357. for {
  358. select {
  359. case <-pb.finish:
  360. return
  361. case <-time.After(pb.RefreshRate):
  362. pb.Update()
  363. }
  364. }
  365. }
  366. type window struct {
  367. Row uint16
  368. Col uint16
  369. Xpixel uint16
  370. Ypixel uint16
  371. }