Merge pull request #11 from nabbar/PkgProgress

Pkg progress

njs-progress/bar.go

@@ -33,6 +33,7 @@ import (
 )
 
 type bar struct {
+	u bool
 	t int64
 	b *mpb.Bar
 	s njs_semaphore.Sem
@@ -41,6 +42,8 @@ type bar struct {
 type Bar interface {
 	Current() int64
 	Completed() bool
+	Increment(n int)
+	Refill(amount int64)
 
 	NewWorker() error
 	NewWorkerTry() bool
@@ -50,16 +53,23 @@ type Bar interface {
 	WaitAll() error
 	Context() context.Context
 	Cancel()
+
+	GetBarMPB() *mpb.Bar
 }
 
-func newBar(b *mpb.Bar, s njs_semaphore.Sem) Bar {
+func newBar(b *mpb.Bar, s njs_semaphore.Sem, total int64) Bar {
 	return &bar{
-		t: 0,
+		u: total > 0,
+		t: total,
 		b: b,
 		s: s,
 	}
 }
 
+func (b bar) GetBarMPB() *mpb.Bar {
+	return b.b
+}
+
 func (b bar) Current() int64 {
 	return b.b.Current()
 }
@@ -68,9 +78,23 @@ func (b bar) Completed() bool {
 	return b.b.Completed()
 }
 
+func (b *bar) Increment(n int) {
+	if n == 0 {
+		n = 1
+	}
+	b.b.IncrBy(n)
+}
+
+func (b *bar) Refill(amount int64) {
+	b.b.SetRefill(amount)
+}
+
 func (b *bar) NewWorker() error {
-	b.t++
-	b.b.SetTotal(b.t, false)
+	if !b.u {
+		b.t++
+		b.b.SetTotal(b.t, false)
+	}
+
 	return b.s.NewWorker()
 }
 

njs-progress/progress.go

@@ -53,6 +53,14 @@ func SetDefaultMessageDone(message string) {
 	defaultMessageDone = message
 }
 
+func GetDefaultStyle() string {
+	return defaultStyle
+}
+
+func GetDefaultMessageDone() string {
+	return defaultMessageDone
+}
+
 type progressBar struct {
 	mpb       *mpb.Progress
 	ctx       context.Context
@@ -62,9 +70,13 @@ type progressBar struct {
 }
 
 type ProgressBar interface {
+	GetMPB() *mpb.Progress
+
 	SetSemaphoreOption(maxSimultaneous int, timeout time.Duration)
-	NewBar(parent context.Context, options ...mpb.BarOption) Bar
-	NewBarSimple(name string) Bar
+
+	NewBar(parent context.Context, total int64, options ...mpb.BarOption) Bar
+	NewBarSimpleETA(name string) Bar
+	NewBarSimpleCounter(name string, total int64) Bar
 }
 
 func NewProgressBar(timeout time.Duration, deadline time.Time, parent context.Context, options ...mpb.ContainerOption) ProgressBar {
@@ -79,23 +91,28 @@ func NewProgressBar(timeout time.Duration, deadline time.Time, parent context.Co
 	}
 }
 
+func (p *progressBar) GetMPB() *mpb.Progress {
+	return p.mpb
+}
+
 func (p *progressBar) SetSemaphoreOption(maxSimultaneous int, timeout time.Duration) {
 	p.sMaxSimul = maxSimultaneous
 	p.sTimeOut = timeout
 }
 
-func (p progressBar) NewBar(parent context.Context, options ...mpb.BarOption) Bar {
+func (p progressBar) NewBar(parent context.Context, total int64, options ...mpb.BarOption) Bar {
 	if parent == nil {
 		parent = p.ctx
 	}
 
 	return newBar(
 		p.mpb.AddBar(0, options...),
 		njs_semaphore.NewSemaphore(p.sMaxSimul, p.sTimeOut, njs_semaphore.GetEmptyTime(), parent),
+		total,
 	)
 }
 
-func (p progressBar) NewBarSimple(name string) Bar {
+func (p progressBar) NewBarSimpleETA(name string) Bar {
 	return newBar(
 		p.mpb.AddBar(0,
 			mpb.BarStyle(defaultStyle),
@@ -110,5 +127,27 @@ func (p progressBar) NewBarSimple(name string) Bar {
 			mpb.AppendDecorators(decor.Percentage()),
 		),
 		njs_semaphore.NewSemaphore(p.sMaxSimul, p.sTimeOut, njs_semaphore.GetEmptyTime(), p.ctx),
+		0,
+	)
+}
+
+func (p progressBar) NewBarSimpleCounter(name string, total int64) Bar {
+	return newBar(
+		p.mpb.AddBar(total,
+			mpb.BarStyle(defaultStyle),
+			mpb.PrependDecorators(
+				// display our name with one space on the right
+				decor.Name(name, decor.WC{W: len(name) + 1, C: decor.DidentRight}),
+				// use counter (no ETA)
+				decor.CountersNoUnit("%d / %d", decor.WCSyncWidth),
+				// replace ETA decorator with "done" message, OnComplete event
+				decor.OnComplete(
+					decor.AverageETA(decor.ET_STYLE_GO, decor.WC{W: 4}), defaultMessageDone,
+				),
+			),
+			mpb.AppendDecorators(decor.Percentage()),
+		),
+		njs_semaphore.NewSemaphore(p.sMaxSimul, p.sTimeOut, njs_semaphore.GetEmptyTime(), p.ctx),
+		total,
 	)
 }