package pb import ( "bytes" "fmt" "io" "os" "strconv" "strings" "sync" "sync/atomic" "text/template" "time" "github.com/mattn/go-colorable" "github.com/mattn/go-isatty" "github.com/cheggaaa/pb/v3/termutil" ) // Version of ProgressBar library const Version = "3.0.5" type key int const ( // Bytes means we're working with byte sizes. Numbers will print as Kb, Mb, etc // bar.Set(pb.Bytes, true) Bytes key = 1 << iota // Use SI bytes prefix names (kB, MB, etc) instead of IEC prefix names (KiB, MiB, etc) SIBytesPrefix // Terminal means we're will print to terminal and can use ascii sequences // Also we're will try to use terminal width Terminal // Static means progress bar will not update automaticly Static // ReturnSymbol - by default in terminal mode it's '\r' ReturnSymbol // Color by default is true when output is tty, but you can set to false for disabling colors Color ) const ( defaultBarWidth = 100 defaultRefreshRate = time.Millisecond * 200 ) // New creates new ProgressBar object func New(total int) *ProgressBar { return New64(int64(total)) } // New64 creates new ProgressBar object using int64 as total func New64(total int64) *ProgressBar { pb := new(ProgressBar) return pb.SetTotal(total) } // StartNew starts new ProgressBar with Default template func StartNew(total int) *ProgressBar { return New(total).Start() } // Start64 starts new ProgressBar with Default template. Using int64 as total. func Start64(total int64) *ProgressBar { return New64(total).Start() } var ( terminalWidth = termutil.TerminalWidth isTerminal = isatty.IsTerminal isCygwinTerminal = isatty.IsCygwinTerminal ) // ProgressBar is the main object of bar type ProgressBar struct { current, total int64 width int maxWidth int mu sync.RWMutex rm sync.Mutex vars map[interface{}]interface{} elements map[string]Element output io.Writer coutput io.Writer nocoutput io.Writer startTime time.Time refreshRate time.Duration tmpl *template.Template state *State buf *bytes.Buffer ticker *time.Ticker finish chan struct{} finished bool configured bool err error } func (pb *ProgressBar) configure() { if pb.configured { return } pb.configured = true if pb.vars == nil { pb.vars = make(map[interface{}]interface{}) } if pb.output == nil { pb.output = os.Stderr } if pb.tmpl == nil { pb.tmpl, pb.err = getTemplate(string(Default)) if pb.err != nil { return } } if pb.vars[Terminal] == nil { if f, ok := pb.output.(*os.File); ok { if isTerminal(f.Fd()) || isCygwinTerminal(f.Fd()) { pb.vars[Terminal] = true } } } if pb.vars[ReturnSymbol] == nil { if tm, ok := pb.vars[Terminal].(bool); ok && tm { pb.vars[ReturnSymbol] = "\r" } } if pb.vars[Color] == nil { if tm, ok := pb.vars[Terminal].(bool); ok && tm { pb.vars[Color] = true } } if pb.refreshRate == 0 { pb.refreshRate = defaultRefreshRate } if f, ok := pb.output.(*os.File); ok { pb.coutput = colorable.NewColorable(f) } else { pb.coutput = pb.output } pb.nocoutput = colorable.NewNonColorable(pb.output) } // Start starts the bar func (pb *ProgressBar) Start() *ProgressBar { pb.mu.Lock() defer pb.mu.Unlock() if pb.finish != nil { return pb } pb.configure() pb.finished = false pb.state = nil pb.startTime = time.Now() if st, ok := pb.vars[Static].(bool); ok && st { return pb } pb.finish = make(chan struct{}) pb.ticker = time.NewTicker(pb.refreshRate) go pb.writer(pb.finish) return pb } func (pb *ProgressBar) writer(finish chan struct{}) { for { select { case <-pb.ticker.C: pb.write(false) case <-finish: pb.ticker.Stop() pb.write(true) finish <- struct{}{} return } } } // Write performs write to the output func (pb *ProgressBar) Write() *ProgressBar { pb.mu.RLock() finished := pb.finished pb.mu.RUnlock() pb.write(finished) return pb } func (pb *ProgressBar) write(finish bool) { result, width := pb.render() if pb.Err() != nil { return } if pb.GetBool(Terminal) { if r := (width - CellCount(result)); r > 0 { result += strings.Repeat(" ", r) } } if ret, ok := pb.Get(ReturnSymbol).(string); ok { result = ret + result if finish && ret == "\r" { result += "\n" } } if pb.GetBool(Color) { pb.coutput.Write([]byte(result)) } else { pb.nocoutput.Write([]byte(result)) } } // Total return current total bar value func (pb *ProgressBar) Total() int64 { return atomic.LoadInt64(&pb.total) } // SetTotal sets the total bar value func (pb *ProgressBar) SetTotal(value int64) *ProgressBar { atomic.StoreInt64(&pb.total, value) return pb } // SetCurrent sets the current bar value func (pb *ProgressBar) SetCurrent(value int64) *ProgressBar { atomic.StoreInt64(&pb.current, value) return pb } // Current return current bar value func (pb *ProgressBar) Current() int64 { return atomic.LoadInt64(&pb.current) } // Add adding given int64 value to bar value func (pb *ProgressBar) Add64(value int64) *ProgressBar { atomic.AddInt64(&pb.current, value) return pb } // Add adding given int value to bar value func (pb *ProgressBar) Add(value int) *ProgressBar { return pb.Add64(int64(value)) } // Increment atomically increments the progress func (pb *ProgressBar) Increment() *ProgressBar { return pb.Add64(1) } // Set sets any value by any key func (pb *ProgressBar) Set(key, value interface{}) *ProgressBar { pb.mu.Lock() defer pb.mu.Unlock() if pb.vars == nil { pb.vars = make(map[interface{}]interface{}) } pb.vars[key] = value return pb } // Get return value by key func (pb *ProgressBar) Get(key interface{}) interface{} { pb.mu.RLock() defer pb.mu.RUnlock() if pb.vars == nil { return nil } return pb.vars[key] } // GetBool return value by key and try to convert there to boolean // If value doesn't set or not boolean - return false func (pb *ProgressBar) GetBool(key interface{}) bool { if v, ok := pb.Get(key).(bool); ok { return v } return false } // SetWidth sets the bar width // When given value <= 0 would be using the terminal width (if possible) or default value. func (pb *ProgressBar) SetWidth(width int) *ProgressBar { pb.mu.Lock() pb.width = width pb.mu.Unlock() return pb } // SetMaxWidth sets the bar maximum width // When given value <= 0 would be using the terminal width (if possible) or default value. func (pb *ProgressBar) SetMaxWidth(maxWidth int) *ProgressBar { pb.mu.Lock() pb.maxWidth = maxWidth pb.mu.Unlock() return pb } // Width return the bar width // It's current terminal width or settled over 'SetWidth' value. func (pb *ProgressBar) Width() (width int) { defer func() { if r := recover(); r != nil { width = defaultBarWidth } }() pb.mu.RLock() width = pb.width maxWidth := pb.maxWidth pb.mu.RUnlock() if width <= 0 { var err error if width, err = terminalWidth(); err != nil { return defaultBarWidth } } if maxWidth > 0 && width > maxWidth { width = maxWidth } return } func (pb *ProgressBar) SetRefreshRate(dur time.Duration) *ProgressBar { pb.mu.Lock() if dur > 0 { pb.refreshRate = dur } pb.mu.Unlock() return pb } // SetWriter sets the io.Writer. Bar will write in this writer // By default this is os.Stderr func (pb *ProgressBar) SetWriter(w io.Writer) *ProgressBar { pb.mu.Lock() pb.output = w pb.configured = false pb.configure() pb.mu.Unlock() return pb } // StartTime return the time when bar started func (pb *ProgressBar) StartTime() time.Time { pb.mu.RLock() defer pb.mu.RUnlock() return pb.startTime } // Format convert int64 to string according to the current settings func (pb *ProgressBar) Format(v int64) string { if pb.GetBool(Bytes) { return formatBytes(v, pb.GetBool(SIBytesPrefix)) } return strconv.FormatInt(v, 10) } // Finish stops the bar func (pb *ProgressBar) Finish() *ProgressBar { pb.mu.Lock() if pb.finished { pb.mu.Unlock() return pb } finishChan := pb.finish pb.finished = true pb.mu.Unlock() if finishChan != nil { finishChan <- struct{}{} <-finishChan pb.mu.Lock() pb.finish = nil pb.mu.Unlock() } return pb } // IsStarted indicates progress bar state func (pb *ProgressBar) IsStarted() bool { pb.mu.RLock() defer pb.mu.RUnlock() return pb.finish != nil } // SetTemplateString sets ProgressBar tempate string and parse it func (pb *ProgressBar) SetTemplateString(tmpl string) *ProgressBar { pb.mu.Lock() defer pb.mu.Unlock() pb.tmpl, pb.err = getTemplate(tmpl) return pb } // SetTemplateString sets ProgressBarTempate and parse it func (pb *ProgressBar) SetTemplate(tmpl ProgressBarTemplate) *ProgressBar { return pb.SetTemplateString(string(tmpl)) } // NewProxyReader creates a wrapper for given reader, but with progress handle // Takes io.Reader or io.ReadCloser // Also, it automatically switches progress bar to handle units as bytes func (pb *ProgressBar) NewProxyReader(r io.Reader) *Reader { pb.Set(Bytes, true) return &Reader{r, pb} } // NewProxyWriter creates a wrapper for given writer, but with progress handle // Takes io.Writer or io.WriteCloser // Also, it automatically switches progress bar to handle units as bytes func (pb *ProgressBar) NewProxyWriter(r io.Writer) *Writer { pb.Set(Bytes, true) return &Writer{r, pb} } func (pb *ProgressBar) render() (result string, width int) { defer func() { if r := recover(); r != nil { pb.SetErr(fmt.Errorf("render panic: %v", r)) } }() pb.rm.Lock() defer pb.rm.Unlock() pb.mu.Lock() pb.configure() if pb.state == nil { pb.state = &State{ProgressBar: pb} pb.buf = bytes.NewBuffer(nil) } if pb.startTime.IsZero() { pb.startTime = time.Now() } pb.state.id++ pb.state.finished = pb.finished pb.state.time = time.Now() pb.mu.Unlock() pb.state.width = pb.Width() width = pb.state.width pb.state.total = pb.Total() pb.state.current = pb.Current() pb.buf.Reset() if e := pb.tmpl.Execute(pb.buf, pb.state); e != nil { pb.SetErr(e) return "", 0 } result = pb.buf.String() aec := len(pb.state.recalc) if aec == 0 { // no adaptive elements return } staticWidth := CellCount(result) - (aec * adElPlaceholderLen) if pb.state.Width()-staticWidth <= 0 { result = strings.Replace(result, adElPlaceholder, "", -1) result = StripString(result, pb.state.Width()) } else { pb.state.adaptiveElWidth = (width - staticWidth) / aec for _, el := range pb.state.recalc { result = strings.Replace(result, adElPlaceholder, el.ProgressElement(pb.state), 1) } } pb.state.recalc = pb.state.recalc[:0] return } // SetErr sets error to the ProgressBar // Error will be available over Err() func (pb *ProgressBar) SetErr(err error) *ProgressBar { pb.mu.Lock() pb.err = err pb.mu.Unlock() return pb } // Err return possible error // When all ok - will be nil // May contain template.Execute errors func (pb *ProgressBar) Err() error { pb.mu.RLock() defer pb.mu.RUnlock() return pb.err } // String return currrent string representation of ProgressBar func (pb *ProgressBar) String() string { res, _ := pb.render() return res } // ProgressElement implements Element interface func (pb *ProgressBar) ProgressElement(s *State, args ...string) string { if s.IsAdaptiveWidth() { pb.SetWidth(s.AdaptiveElWidth()) } return pb.String() } // State represents the current state of bar // Need for bar elements type State struct { *ProgressBar id uint64 total, current int64 width, adaptiveElWidth int finished, adaptive bool time time.Time recalc []Element } // Id it's the current state identifier // - incremental // - starts with 1 // - resets after finish/start func (s *State) Id() uint64 { return s.id } // Total it's bar int64 total func (s *State) Total() int64 { return s.total } // Value it's current value func (s *State) Value() int64 { return s.current } // Width of bar func (s *State) Width() int { return s.width } // AdaptiveElWidth - adaptive elements must return string with given cell count (when AdaptiveElWidth > 0) func (s *State) AdaptiveElWidth() int { return s.adaptiveElWidth } // IsAdaptiveWidth returns true when element must be shown as adaptive func (s *State) IsAdaptiveWidth() bool { return s.adaptive } // IsFinished return true when bar is finished func (s *State) IsFinished() bool { return s.finished } // IsFirst return true only in first render func (s *State) IsFirst() bool { return s.id == 1 } // Time when state was created func (s *State) Time() time.Time { return s.time }