diff options
Diffstat (limited to 'vendor/github.com/vbauerster/mpb/v5/bar.go')
-rw-r--r-- | vendor/github.com/vbauerster/mpb/v5/bar.go | 476 |
1 files changed, 476 insertions, 0 deletions
diff --git a/vendor/github.com/vbauerster/mpb/v5/bar.go b/vendor/github.com/vbauerster/mpb/v5/bar.go new file mode 100644 index 000000000..1a4c66fe1 --- /dev/null +++ b/vendor/github.com/vbauerster/mpb/v5/bar.go @@ -0,0 +1,476 @@ +package mpb + +import ( + "bytes" + "context" + "fmt" + "io" + "log" + "strings" + "time" + "unicode/utf8" + + "github.com/vbauerster/mpb/v5/decor" +) + +// BarFiller interface. +// Bar renders itself by calling BarFiller's Fill method. You can +// literally have any bar kind, by implementing this interface and +// passing it to the *Progress.Add(...) *Bar method. +type BarFiller interface { + Fill(w io.Writer, width int, stat *decor.Statistics) +} + +// BarFillerFunc is function type adapter to convert function into Filler. +type BarFillerFunc func(w io.Writer, width int, stat *decor.Statistics) + +func (f BarFillerFunc) Fill(w io.Writer, width int, stat *decor.Statistics) { + f(w, width, stat) +} + +// Bar represents a progress Bar. +type Bar struct { + priority int // used by heap + index int // used by heap + + extendedLines int + toShutdown bool + toDrop bool + noPop bool + hasEwmaDecorators bool + operateState chan func(*bState) + frameCh chan io.Reader + syncTableCh chan [][]chan int + completed chan bool + + // cancel is called either by user or on complete event + cancel func() + // done is closed after cacheState is assigned + done chan struct{} + // cacheState is populated, right after close(shutdown) + cacheState *bState + + container *Progress + dlogger *log.Logger + recoveredPanic interface{} +} + +type extFunc func(in io.Reader, tw int, st *decor.Statistics) (out io.Reader, lines int) + +type bState struct { + baseF BarFiller + filler BarFiller + id int + width int + total int64 + current int64 + lastN int64 + iterated bool + trimSpace bool + toComplete bool + completeFlushed bool + noPop bool + aDecorators []decor.Decorator + pDecorators []decor.Decorator + averageDecorators []decor.AverageDecorator + ewmaDecorators []decor.EwmaDecorator + shutdownListeners []decor.ShutdownListener + bufP, bufB, bufA *bytes.Buffer + extender extFunc + + // priority overrides *Bar's priority, if set + priority int + // dropOnComplete propagates to *Bar + dropOnComplete bool + // runningBar is a key for *pState.parkedBars + runningBar *Bar + + debugOut io.Writer +} + +func newBar(container *Progress, bs *bState) *Bar { + logPrefix := fmt.Sprintf("%sbar#%02d ", container.dlogger.Prefix(), bs.id) + ctx, cancel := context.WithCancel(container.ctx) + + bar := &Bar{ + container: container, + priority: bs.priority, + toDrop: bs.dropOnComplete, + noPop: bs.noPop, + operateState: make(chan func(*bState)), + frameCh: make(chan io.Reader, 1), + syncTableCh: make(chan [][]chan int), + completed: make(chan bool, 1), + done: make(chan struct{}), + cancel: cancel, + dlogger: log.New(bs.debugOut, logPrefix, log.Lshortfile), + } + + go bar.serve(ctx, bs) + return bar +} + +// ProxyReader wraps r with metrics required for progress tracking. +// Panics if r is nil. +func (b *Bar) ProxyReader(r io.Reader) io.ReadCloser { + if r == nil { + panic("expected non nil io.Reader") + } + return newProxyReader(r, b) +} + +// ID returs id of the bar. +func (b *Bar) ID() int { + result := make(chan int) + select { + case b.operateState <- func(s *bState) { result <- s.id }: + return <-result + case <-b.done: + return b.cacheState.id + } +} + +// Current returns bar's current number, in other words sum of all increments. +func (b *Bar) Current() int64 { + result := make(chan int64) + select { + case b.operateState <- func(s *bState) { result <- s.current }: + return <-result + case <-b.done: + return b.cacheState.current + } +} + +// SetRefill fills bar with refill rune up to amount argument. +// Given default bar style is "[=>-]<+", refill rune is '+'. +// To set bar style use mpb.BarStyle(string) BarOption. +func (b *Bar) SetRefill(amount int64) { + type refiller interface { + SetRefill(int64) + } + b.operateState <- func(s *bState) { + if f, ok := s.baseF.(refiller); ok { + f.SetRefill(amount) + } + } +} + +// TraverseDecorators traverses all available decorators and calls cb func on each. +func (b *Bar) TraverseDecorators(cb func(decor.Decorator)) { + b.operateState <- func(s *bState) { + for _, decorators := range [...][]decor.Decorator{ + s.pDecorators, + s.aDecorators, + } { + for _, d := range decorators { + cb(extractBaseDecorator(d)) + } + } + } +} + +// SetTotal sets total dynamically. +// If total is less or equal to zero it takes progress' current value. +// If complete is true, complete event will be triggered. +func (b *Bar) SetTotal(total int64, complete bool) { + select { + case b.operateState <- func(s *bState) { + if total <= 0 { + s.total = s.current + } else { + s.total = total + } + if complete && !s.toComplete { + s.current = s.total + s.toComplete = true + go b.refreshTillShutdown() + } + }: + case <-b.done: + } +} + +// SetCurrent sets progress' current to an arbitrary value. +func (b *Bar) SetCurrent(current int64) { + select { + case b.operateState <- func(s *bState) { + s.iterated = true + s.lastN = current - s.current + s.current = current + if s.total > 0 && s.current >= s.total { + s.current = s.total + s.toComplete = true + go b.refreshTillShutdown() + } + }: + case <-b.done: + } +} + +// Increment is a shorthand for b.IncrInt64(1). +func (b *Bar) Increment() { + b.IncrInt64(1) +} + +// IncrBy is a shorthand for b.IncrInt64(int64(n)). +func (b *Bar) IncrBy(n int) { + b.IncrInt64(int64(n)) +} + +// IncrInt64 increments progress by amount of n. +func (b *Bar) IncrInt64(n int64) { + select { + case b.operateState <- func(s *bState) { + s.iterated = true + s.lastN = n + s.current += n + if s.total > 0 && s.current >= s.total { + s.current = s.total + s.toComplete = true + go b.refreshTillShutdown() + } + }: + case <-b.done: + } +} + +// DecoratorEwmaUpdate updates all EWMA based decorators. Should be +// called on each iteration, because EWMA's unit of measure is an +// iteration's duration. Panics if called before *Bar.Incr... family +// methods. +func (b *Bar) DecoratorEwmaUpdate(dur time.Duration) { + select { + case b.operateState <- func(s *bState) { + ewmaIterationUpdate(false, s, dur) + }: + case <-b.done: + ewmaIterationUpdate(true, b.cacheState, dur) + } +} + +// DecoratorAverageAdjust adjusts all average based decorators. Call +// if you need to adjust start time of all average based decorators +// or after progress resume. +func (b *Bar) DecoratorAverageAdjust(start time.Time) { + select { + case b.operateState <- func(s *bState) { + for _, d := range s.averageDecorators { + d.AverageAdjust(start) + } + }: + case <-b.done: + } +} + +// SetPriority changes bar's order among multiple bars. Zero is highest +// priority, i.e. bar will be on top. If you don't need to set priority +// dynamically, better use BarPriority option. +func (b *Bar) SetPriority(priority int) { + select { + case <-b.done: + default: + b.container.setBarPriority(b, priority) + } +} + +// Abort interrupts bar's running goroutine. Call this, if you'd like +// to stop/remove bar before completion event. It has no effect after +// completion event. If drop is true bar will be removed as well. +func (b *Bar) Abort(drop bool) { + select { + case <-b.done: + default: + if drop { + b.container.dropBar(b) + } + b.cancel() + } +} + +// Completed reports whether the bar is in completed state. +func (b *Bar) Completed() bool { + select { + case b.operateState <- func(s *bState) { b.completed <- s.toComplete }: + return <-b.completed + case <-b.done: + return true + } +} + +func (b *Bar) serve(ctx context.Context, s *bState) { + defer b.container.bwg.Done() + for { + select { + case op := <-b.operateState: + op(s) + case <-ctx.Done(): + b.cacheState = s + close(b.done) + // Notifying decorators about shutdown event + for _, sl := range s.shutdownListeners { + sl.Shutdown() + } + return + } + } +} + +func (b *Bar) render(tw int) { + if b.recoveredPanic != nil { + b.toShutdown = false + b.frameCh <- b.panicToFrame(tw) + return + } + select { + case b.operateState <- func(s *bState) { + defer func() { + // recovering if user defined decorator panics for example + if p := recover(); p != nil { + b.dlogger.Println(p) + b.recoveredPanic = p + b.toShutdown = !s.completeFlushed + b.frameCh <- b.panicToFrame(tw) + } + }() + + st := newStatistics(s) + frame := s.draw(tw, st) + frame, b.extendedLines = s.extender(frame, tw, st) + + b.toShutdown = s.toComplete && !s.completeFlushed + s.completeFlushed = s.toComplete + b.frameCh <- frame + }: + case <-b.done: + s := b.cacheState + st := newStatistics(s) + frame := s.draw(tw, st) + frame, b.extendedLines = s.extender(frame, tw, st) + b.frameCh <- frame + } +} + +func (b *Bar) panicToFrame(termWidth int) io.Reader { + return strings.NewReader(fmt.Sprintf(fmt.Sprintf("%%.%dv\n", termWidth), b.recoveredPanic)) +} + +func (b *Bar) subscribeDecorators() { + var averageDecorators []decor.AverageDecorator + var ewmaDecorators []decor.EwmaDecorator + var shutdownListeners []decor.ShutdownListener + b.TraverseDecorators(func(d decor.Decorator) { + if d, ok := d.(decor.AverageDecorator); ok { + averageDecorators = append(averageDecorators, d) + } + if d, ok := d.(decor.EwmaDecorator); ok { + ewmaDecorators = append(ewmaDecorators, d) + } + if d, ok := d.(decor.ShutdownListener); ok { + shutdownListeners = append(shutdownListeners, d) + } + }) + b.operateState <- func(s *bState) { + s.averageDecorators = averageDecorators + s.ewmaDecorators = ewmaDecorators + s.shutdownListeners = shutdownListeners + } + b.hasEwmaDecorators = len(ewmaDecorators) != 0 +} + +func (b *Bar) refreshTillShutdown() { + for { + select { + case b.container.refreshCh <- time.Now(): + case <-b.done: + return + } + } +} + +func (b *Bar) wSyncTable() [][]chan int { + select { + case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }: + return <-b.syncTableCh + case <-b.done: + return b.cacheState.wSyncTable() + } +} + +func (s *bState) draw(termWidth int, stat *decor.Statistics) io.Reader { + for _, d := range s.pDecorators { + s.bufP.WriteString(d.Decor(stat)) + } + + for _, d := range s.aDecorators { + s.bufA.WriteString(d.Decor(stat)) + } + + s.bufA.WriteByte('\n') + + prependCount := utf8.RuneCount(s.bufP.Bytes()) + appendCount := utf8.RuneCount(s.bufA.Bytes()) - 1 + + if fitWidth := s.width; termWidth > 1 { + if !s.trimSpace { + // reserve space for edge spaces + termWidth -= 2 + s.bufB.WriteByte(' ') + defer s.bufB.WriteByte(' ') + } + if prependCount+s.width+appendCount > termWidth { + fitWidth = termWidth - prependCount - appendCount + } + s.filler.Fill(s.bufB, fitWidth, stat) + } + + return io.MultiReader(s.bufP, s.bufB, s.bufA) +} + +func (s *bState) wSyncTable() [][]chan int { + columns := make([]chan int, 0, len(s.pDecorators)+len(s.aDecorators)) + var pCount int + for _, d := range s.pDecorators { + if ch, ok := d.Sync(); ok { + columns = append(columns, ch) + pCount++ + } + } + var aCount int + for _, d := range s.aDecorators { + if ch, ok := d.Sync(); ok { + columns = append(columns, ch) + aCount++ + } + } + table := make([][]chan int, 2) + table[0] = columns[0:pCount] + table[1] = columns[pCount : pCount+aCount : pCount+aCount] + return table +} + +func newStatistics(s *bState) *decor.Statistics { + return &decor.Statistics{ + ID: s.id, + Completed: s.completeFlushed, + Total: s.total, + Current: s.current, + } +} + +func extractBaseDecorator(d decor.Decorator) decor.Decorator { + if d, ok := d.(decor.Wrapper); ok { + return extractBaseDecorator(d.Base()) + } + return d +} + +func ewmaIterationUpdate(done bool, s *bState, dur time.Duration) { + if !done && !s.iterated { + panic("increment required before ewma iteration update") + } else { + s.iterated = false + } + for _, d := range s.ewmaDecorators { + d.EwmaUpdate(s.lastN, dur) + } +} |