这是indexloc提供的服务,不要输入任何密码
Skip to content
54 changes: 37 additions & 17 deletions progress.go
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,10 @@ type Progress struct {
// RefreshInterval in the time duration to wait for refreshing the output
RefreshInterval time.Duration

lw *uilive.Writer
stopChan chan struct{}
mtx *sync.RWMutex
lw *uilive.Writer
ticker *time.Ticker
tdone chan bool
mtx *sync.RWMutex
}

// New returns a new progress bar with defaults
Expand All @@ -46,9 +47,8 @@ func New() *Progress {
Bars: make([]*Bar, 0),
RefreshInterval: RefreshInterval,

lw: uilive.New(),
stopChan: make(chan struct{}),
mtx: &sync.RWMutex{},
lw: uilive.New(),
mtx: &sync.RWMutex{},
}
}

Expand Down Expand Up @@ -85,35 +85,55 @@ func (p *Progress) AddBar(total int) *Bar {

// Listen listens for updates and renders the progress bars
func (p *Progress) Listen() {
var tickChan = p.ticker.C
p.lw.Out = p.Out

for {
select {
case <-p.stopChan:
return
default:
time.Sleep(p.RefreshInterval)
case <-tickChan:
p.mtx.RLock()
for _, bar := range p.Bars {
fmt.Fprintln(p.lw, bar.String())

if p.ticker != nil {
p.print()
p.lw.Flush()
}
p.lw.Flush()

p.mtx.RUnlock()
case <-p.tdone:
if p.ticker != nil {
p.ticker.Stop()
p.ticker = nil
return
}
}
}
}

func (p *Progress) print() {
for _, bar := range p.Bars {
fmt.Fprintln(p.lw, bar.String())
}
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This line would allow earlier to restart the listener. Suggested implementation doesn't support this anymore. Not saying whether this is important or not, just a word of caution.


// Start starts the rendering the progress of progress bars. It listens for updates using `bar.Set(n)` and new bars when added using `AddBar`
func (p *Progress) Start() {
if p.stopChan == nil {
p.stopChan = make(chan struct{})
p.mtx.Lock()
if p.ticker == nil {
p.ticker = time.NewTicker(RefreshInterval)
p.tdone = make(chan bool, 1)
}
p.mtx.Unlock()

go p.Listen()
}

// Stop stops listening
func (p *Progress) Stop() {
close(p.stopChan)
p.stopChan = nil
p.mtx.Lock()
close(p.tdone)
p.print()
p.lw.Flush()
p.mtx.Unlock()
}

// Bypass returns a writer which allows non-buffered data to be written to the underlying output
Expand Down
43 changes: 43 additions & 0 deletions progress_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
package uiprogress

import (
"bytes"
"fmt"
"strings"
"sync"
"testing"
"time"
)

func TestStoppingPrintout(t *testing.T) {
progress := New()
progress.RefreshInterval = time.Millisecond * 10
var buffer = &bytes.Buffer{}
progress.Out = buffer
bar := progress.AddBar(100)
progress.Start()

var wg sync.WaitGroup

wg.Add(1)

go func() {
for i := 0; i <= 80; i = i + 10 {
bar.Set(i)
time.Sleep(time.Millisecond * 5)
}

wg.Done()
}()

wg.Wait()

progress.Stop()
fmt.Fprintf(buffer, "foo")

var wantSuffix = "[======================================================>-------------]\nfoo"

if !strings.HasSuffix(buffer.String(), wantSuffix) {
t.Errorf("Content that should be printed after stop not appearing on buffer.")
}
}