github.com/emc-advanced-dev/unik@v0.0.0-20190717152701-a58d3e8e33b7/pkg/util/download.go (about)

     1  package util
     2  
     3  import (
     4  	"gopkg.in/cheggaaa/pb.v1"
     5  	"io"
     6  )
     7  
     8  ///http://stackoverflow.com/questions/22421375/how-to-print-the-bytes-while-the-file-is-being-downloaded-golang
     9  // WriteCounter counts the number of bytes written to it.
    10  type writeCounter struct {
    11  	current int64 // Total # of bytes transferred
    12  	total   int64 // Expected length
    13  	bar     *pb.ProgressBar
    14  }
    15  
    16  func newWriteCounter(total int64) *writeCounter {
    17  	return &writeCounter{
    18  		total: total,
    19  		bar:   pb.StartNew(int(total)),
    20  	}
    21  }
    22  
    23  // Write implements the io.Writer interface.
    24  //
    25  // Always completes and never returns an error.
    26  func (wc *writeCounter) Write(p []byte) (int, error) {
    27  	n := len(p)
    28  	wc.current += int64(n)
    29  	wc.bar.Set(int(wc.current))
    30  	if wc.current >= wc.total-1 {
    31  		wc.bar.FinishPrint("download complete")
    32  	}
    33  	return n, nil
    34  }
    35  
    36  func ReaderWithProgress(r io.Reader, total int64) io.Reader {
    37  	return io.TeeReader(r, newWriteCounter(total))
    38  }