github.com/benchkram/bob@v0.0.0-20240314204020-b7a57f2f9be9/pkg/progress/progress.go (about)

     1  package progress
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"math"
     7  	"os"
     8  	"sync"
     9  	"time"
    10  )
    11  
    12  // Progress tracks progress rendering at a certain interval on a new line
    13  // the current percentage of bytes added from a total of maxBytes
    14  //
    15  // Inspired by https://github.com/schollz/progressbar
    16  type Progress struct {
    17  	// maxBytes is the total size of bytes that are tracked
    18  	maxBytes int64
    19  	// currentBytes is the total number of bytes tracked so far
    20  	currentBytes int64
    21  
    22  	// description is added as prefix on every render
    23  	description string
    24  
    25  	// currentPercent is percent of currentBytes from maxBytes
    26  	currentPercent int
    27  	// lastPercent is the last rendered percent
    28  	lastPercent int
    29  
    30  	// lock for Add operations
    31  	lock sync.Mutex
    32  
    33  	// lastRendered time
    34  	lastRendered time.Time
    35  	// intervalToRender sets an interval when to render the progress
    36  	intervalToRender time.Duration
    37  }
    38  
    39  // NewProgress initialize a new progress ready to use
    40  func NewProgress(maxBytes int64, description string, intervalToShow time.Duration) *Progress {
    41  	return &Progress{
    42  		maxBytes:         maxBytes,
    43  		description:      description,
    44  		lastRendered:     time.Now(),
    45  		intervalToRender: intervalToShow,
    46  	}
    47  }
    48  
    49  // Add will add the specified amount to the progressbar
    50  func (p *Progress) Add(num int) {
    51  	p.Add64(int64(num))
    52  }
    53  
    54  // Add64 will add the specified amount to the progressbar
    55  func (p *Progress) Add64(num int64) {
    56  	p.lock.Lock()
    57  	defer p.lock.Unlock()
    58  
    59  	p.currentBytes += num
    60  	p.currentPercent = int(float64(p.currentBytes) / float64(p.maxBytes) * 100)
    61  
    62  	if p.currentPercent == p.lastPercent {
    63  		return
    64  	}
    65  
    66  	if p.currentBytes == p.maxBytes || time.Since(p.lastRendered) >= p.intervalToRender {
    67  		p.render()
    68  	}
    69  }
    70  
    71  // render current progress ex. `description 54% (7.4kB/7.4kB)`
    72  func (p *Progress) render() {
    73  	currentHuman, currentUnit := humanizeBytes(float64(p.currentBytes))
    74  	maxHuman, maxUnit := humanizeBytes(float64(p.maxBytes))
    75  
    76  	fmt.Fprintf(os.Stdout, "%s %d%% (%s%s/%s%s)\n", p.description, p.currentPercent, currentHuman, currentUnit, maxHuman, maxUnit)
    77  	p.lastRendered = time.Now()
    78  	p.lastPercent = p.currentPercent
    79  }
    80  
    81  // Finish sets the current progress to 100%
    82  func (p *Progress) Finish() {
    83  	p.lock.Lock()
    84  	p.currentBytes = p.maxBytes
    85  	p.lock.Unlock()
    86  	p.Add(0)
    87  }
    88  
    89  // Reader will wrap an io.Reader adding progress functionality
    90  type Reader struct {
    91  	io.Reader
    92  	bar *Progress
    93  }
    94  
    95  // NewReader return a new Reader with a given progress
    96  func NewReader(r io.Reader, bar *Progress) Reader {
    97  	return Reader{
    98  		Reader: r,
    99  		bar:    bar,
   100  	}
   101  }
   102  
   103  // Read will read n bytes and track progress
   104  func (r *Reader) Read(p []byte) (n int, err error) {
   105  	n, err = r.Reader.Read(p)
   106  	r.bar.Add(n)
   107  	return n, err
   108  }
   109  
   110  // Close the reader when it implements io.Closer
   111  func (r *Reader) Close() (err error) {
   112  	if closer, ok := r.Reader.(io.Closer); ok {
   113  		return closer.Close()
   114  	}
   115  	r.bar.Finish()
   116  	return
   117  }
   118  
   119  func humanizeBytes(s float64) (string, string) {
   120  	sizes := []string{"B", "kB", "MB", "GB", "TB", "PB", "EB"}
   121  	base := 1024.0
   122  	if s < 10 {
   123  		return fmt.Sprintf("%2.0f", s), sizes[0]
   124  	}
   125  	e := math.Floor(logn(s, base))
   126  	suffix := sizes[int(e)]
   127  	val := math.Floor(s/math.Pow(base, e)*10+0.5) / 10
   128  	f := "%.0f"
   129  	if val < 10 {
   130  		f = "%.1f"
   131  	}
   132  
   133  	return fmt.Sprintf(f, val), suffix
   134  }
   135  
   136  func logn(n, b float64) float64 {
   137  	return math.Log(n) / math.Log(b)
   138  }