progressreader.go 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. package utils
  2. import (
  3. "io"
  4. "time"
  5. )
  6. // Reader with progress bar
  7. type progressReader struct {
  8. reader io.ReadCloser // Stream to read from
  9. output io.Writer // Where to send progress bar to
  10. progress JSONProgress
  11. lastUpdate int // How many bytes read at least update
  12. ID string
  13. action string
  14. sf *StreamFormatter
  15. newLine bool
  16. }
  17. func (r *progressReader) Read(p []byte) (n int, err error) {
  18. read, err := r.reader.Read(p)
  19. r.progress.Current += read
  20. updateEvery := 1024 * 512 //512kB
  21. if r.progress.Total > 0 {
  22. // Update progress for every 1% read if 1% < 512kB
  23. if increment := int(0.01 * float64(r.progress.Total)); increment < updateEvery {
  24. updateEvery = increment
  25. }
  26. }
  27. if r.progress.Current-r.lastUpdate > updateEvery || err != nil {
  28. r.output.Write(r.sf.FormatProgress(r.ID, r.action, &r.progress))
  29. r.lastUpdate = r.progress.Current
  30. }
  31. // Send newline when complete
  32. if r.newLine && err != nil && read == 0 {
  33. r.output.Write(r.sf.FormatStatus("", ""))
  34. }
  35. return read, err
  36. }
  37. func (r *progressReader) Close() error {
  38. r.progress.Current = r.progress.Total
  39. r.output.Write(r.sf.FormatProgress(r.ID, r.action, &r.progress))
  40. return r.reader.Close()
  41. }
  42. func ProgressReader(r io.ReadCloser, size int, output io.Writer, sf *StreamFormatter, newline bool, ID, action string) *progressReader {
  43. return &progressReader{
  44. reader: r,
  45. output: NewWriteFlusher(output),
  46. ID: ID,
  47. action: action,
  48. progress: JSONProgress{Total: size, Start: time.Now().UTC().Unix()},
  49. sf: sf,
  50. newLine: newline,
  51. }
  52. }