Merge pull request #188 from cheggaaa/v3-pooling

multiple progress bars on v3

v3/go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/mattn/go-isatty v0.0.12
 	github.com/mattn/go-runewidth v0.0.12
 	github.com/rivo/uniseg v0.2.0 // indirect
-	golang.org/x/sys v0.0.0-20210403161142-5e06dd20ab57 // indirect
+	golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6 // indirect
 )
 
 go 1.12

v3/go.sum

@@ -13,5 +13,5 @@ github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200223170610-d5e6a3e2c0ae/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210403161142-5e06dd20ab57 h1:F5Gozwx4I1xtr/sr/8CFbb57iKi3297KFs0QDbGN60A=
-golang.org/x/sys v0.0.0-20210403161142-5e06dd20ab57/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6 h1:nonptSpoQ4vQjyraW20DXPAglgQfVnM9ZC6MmNLMR60=
+golang.org/x/sys v0.0.0-20220503163025-988cb79eb6c6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

v3/pb.go

@@ -408,6 +408,13 @@ func (pb *ProgressBar) IsStarted() bool {
 	return pb.finish != nil
 }
 
+// IsFinished indicates progress bar is finished
+func (pb *ProgressBar) IsFinished() bool {
+	pb.mu.RLock()
+	defer pb.mu.RUnlock()
+	return pb.finished
+}
+
 // SetTemplateString sets ProgressBar tempate string and parse it
 func (pb *ProgressBar) SetTemplateString(tmpl string) *ProgressBar {
 	pb.mu.Lock()

v3/pool.go

@@ -0,0 +1,105 @@
+// +build linux darwin freebsd netbsd openbsd solaris dragonfly windows plan9 aix
+
+package pb
+
+import (
+	"io"
+	"sync"
+	"time"
+
+	"github.com/cheggaaa/pb/v3/termutil"
+)
+
+// Create and start new pool with given bars
+// You need call pool.Stop() after work
+func StartPool(pbs ...*ProgressBar) (pool *Pool, err error) {
+	pool = new(Pool)
+	if err = pool.Start(); err != nil {
+		return
+	}
+	pool.Add(pbs...)
+	return
+}
+
+// NewPool initialises a pool with progress bars, but
+// doesn't start it. You need to call Start manually
+func NewPool(pbs ...*ProgressBar) (pool *Pool) {
+	pool = new(Pool)
+	pool.Add(pbs...)
+	return
+}
+
+type Pool struct {
+	Output        io.Writer
+	RefreshRate   time.Duration
+	bars          []*ProgressBar
+	lastBarsCount int
+	shutdownCh    chan struct{}
+	workerCh      chan struct{}
+	m             sync.Mutex
+	finishOnce    sync.Once
+}
+
+// Add progress bars.
+func (p *Pool) Add(pbs ...*ProgressBar) {
+	p.m.Lock()
+	defer p.m.Unlock()
+	for _, bar := range pbs {
+		bar.Set(Static, true)
+		bar.Start()
+		p.bars = append(p.bars, bar)
+	}
+}
+
+func (p *Pool) Start() (err error) {
+	p.RefreshRate = defaultRefreshRate
+	p.shutdownCh, err = termutil.RawModeOn()
+	if err != nil {
+		return
+	}
+	p.workerCh = make(chan struct{})
+	go p.writer()
+	return
+}
+
+func (p *Pool) writer() {
+	var first = true
+	defer func() {
+		if first == false {
+			p.print(false)
+		} else {
+			p.print(true)
+			p.print(false)
+		}
+		close(p.workerCh)
+	}()
+
+	for {
+		select {
+		case <-time.After(p.RefreshRate):
+			if p.print(first) {
+				p.print(false)
+				return
+			}
+			first = false
+		case <-p.shutdownCh:
+			return
+		}
+	}
+}
+
+// Restore terminal state and close pool
+func (p *Pool) Stop() error {
+	p.finishOnce.Do(func() {
+		if p.shutdownCh != nil {
+			close(p.shutdownCh)
+		}
+	})
+
+	// Wait for the worker to complete
+	select {
+	case <-p.workerCh:
+	}
+
+	return termutil.RawModeOff()
+}

v3/pool_win.go

@@ -0,0 +1,46 @@
+// +build windows
+
+package pb
+
+import (
+	"fmt"
+	"log"
+
+	"github.com/cheggaaa/pb/v3/termutil"
+)
+
+func (p *Pool) print(first bool) bool {
+	p.m.Lock()
+	defer p.m.Unlock()
+	var out string
+	if !first {
+		coords, err := termutil.GetCursorPos()
+		if err != nil {
+			log.Panic(err)
+		}
+		coords.Y -= int16(p.lastBarsCount)
+		if coords.Y < 0 {
+			coords.Y = 0
+		}
+		coords.X = 0
+
+		err =  termutil.SetCursorPos(coords)
+		if err != nil {
+			log.Panic(err)
+		}
+	}
+	isFinished := true
+	for _, bar := range p.bars {
+		if !bar.IsFinished() {
+			isFinished = false
+		}
+		out += fmt.Sprintf("\r%s\n", bar.String())
+	}
+	if p.Output != nil {
+		fmt.Fprint(p.Output, out)
+	} else {
+		fmt.Print(out)
+	}
+	p.lastBarsCount = len(p.bars)
+	return isFinished
+}

v3/pool_x.go

@@ -0,0 +1,43 @@
+// +build linux darwin freebsd netbsd openbsd solaris dragonfly plan9 aix
+
+package pb
+
+import (
+	"fmt"
+	"os"
+
+	"github.com/cheggaaa/pb/v3/termutil"
+)
+
+func (p *Pool) print(first bool) bool {
+	p.m.Lock()
+	defer p.m.Unlock()
+	var out string
+	if !first {
+		out = fmt.Sprintf("\033[%dA", p.lastBarsCount)
+	}
+	isFinished := true
+	bars := p.bars
+	rows, cols, err := termutil.TerminalSize()
+	if err != nil {
+		cols = defaultBarWidth
+	}
+	if rows > 0 && len(bars) > rows {
+		// we need to hide bars that overflow terminal height
+		bars = bars[len(bars)-rows:]
+	}
+	for _, bar := range bars {
+		if !bar.IsFinished() {
+			isFinished = false
+		}
+		bar.SetWidth(cols)
+		out += fmt.Sprintf("\r%s\n", bar.String())
+	}
+	if p.Output != nil {
+		fmt.Fprint(p.Output, out)
+	} else {
+		fmt.Fprint(os.Stderr, out)
+	}
+	p.lastBarsCount = len(bars)
+	return isFinished
+}

v3/termutil/term.go

@@ -10,6 +10,16 @@ import (
 var echoLocked bool
 var echoLockMutex sync.Mutex
 var errLocked = errors.New("terminal locked")
+var autoTerminate = true
+
+// AutoTerminate enables or disables automatic terminate signal catching.
+// It's needed to restore the terminal state after the pool was used.
+// By default, it's enabled.
+func AutoTerminate(enable bool) {
+	echoLockMutex.Lock()
+	defer echoLockMutex.Unlock()
+	autoTerminate = enable
+}
 
 // RawModeOn switches terminal to raw mode
 func RawModeOn() (quit chan struct{}, err error) {
@@ -45,8 +55,10 @@ func RawModeOff() (err error) {
 // listen exit signals and restore terminal state
 func catchTerminate(quit chan struct{}) {
 	sig := make(chan os.Signal, 1)
-	signal.Notify(sig, unlockSignals...)
-	defer signal.Stop(sig)
+	if autoTerminate {
+		signal.Notify(sig, unlockSignals...)
+		defer signal.Stop(sig)
+	}
 	select {
 	case <-quit:
 		RawModeOff()

v3/termutil/term_appengine.go

@@ -1,3 +1,4 @@
+//go:build appengine
 // +build appengine
 
 package termutil

v3/termutil/term_bsd.go

@@ -1,3 +1,4 @@
+//go:build (darwin || freebsd || netbsd || openbsd || dragonfly) && !appengine
 // +build darwin freebsd netbsd openbsd dragonfly
 // +build !appengine
 

v3/termutil/term_linux.go

@@ -1,5 +1,5 @@
-// +build linux
-// +build !appengine
+//go:build linux && !appengine
+// +build linux,!appengine
 
 package termutil
 

v3/termutil/term_nix.go

@@ -1,3 +1,4 @@
+//go:build (linux || darwin || freebsd || netbsd || openbsd || dragonfly) && !appengine
 // +build linux darwin freebsd netbsd openbsd dragonfly
 // +build !appengine
 

v3/termutil/term_solaris.go

@@ -1,5 +1,5 @@
-// +build solaris
-// +build !appengine
+//go:build solaris && !appengine
+// +build solaris,!appengine
 
 package termutil
 

v3/termutil/term_win.go

@@ -1,3 +1,4 @@
+//go:build windows
 // +build windows
 
 package termutil
@@ -111,7 +112,7 @@ func termWidthTPut() (width int, err error) {
 	return strconv.Atoi(string(res))
 }
 
-func getCursorPos() (pos coordinates, err error) {
+func GetCursorPos() (pos coordinates, err error) {
 	var info consoleScreenBufferInfo
 	_, _, e := syscall.Syscall(procGetConsoleScreenBufferInfo.Addr(), 2, uintptr(syscall.Stdout), uintptr(unsafe.Pointer(&info)), 0)
 	if e != 0 {
@@ -120,7 +121,7 @@ func getCursorPos() (pos coordinates, err error) {
 	return info.dwCursorPosition, nil
 }
 
-func setCursorPos(pos coordinates) error {
+func SetCursorPos(pos coordinates) error {
 	_, _, e := syscall.Syscall(setConsoleCursorPosition.Addr(), 2, uintptr(syscall.Stdout), uintptr(uint32(uint16(pos.Y))<<16|uint32(uint16(pos.X))), 0)
 	if e != 0 {
 		return error(e)

v3/termutil/term_x.go

@@ -1,3 +1,4 @@
+//go:build (linux || darwin || freebsd || netbsd || openbsd || solaris || dragonfly) && !appengine
 // +build linux darwin freebsd netbsd openbsd solaris dragonfly
 // +build !appengine
 
@@ -16,6 +17,7 @@ var (
 	unlockSignals = []os.Signal{
 		os.Interrupt, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGKILL,
 	}
+	oldState syscall.Termios
 )
 
 type window struct {
@@ -35,42 +37,45 @@ func init() {
 
 // TerminalWidth returns width of the terminal.
 func TerminalWidth() (int, error) {
+	_, c, err := TerminalSize()
+	return c, err
+}
+
+// TerminalSize returns size of the terminal.
+func TerminalSize() (rows, cols int, err error) {
 	w := new(window)
 	res, _, err := syscall.Syscall(sysIoctl,
 		tty.Fd(),
 		uintptr(syscall.TIOCGWINSZ),
 		uintptr(unsafe.Pointer(w)),
 	)
 	if int(res) == -1 {
-		return 0, err
+		return 0, 0, err
 	}
-	return int(w.Col), nil
+	return int(w.Row), int(w.Col), nil
 }
 
-var oldState syscall.Termios
-
-func lockEcho() (err error) {
+func lockEcho() error {
 	fd := tty.Fd()
-	if _, _, e := syscall.Syscall6(sysIoctl, fd, ioctlReadTermios, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0); e != 0 {
-		err = fmt.Errorf("Can't get terminal settings: %v", e)
-		return
+
+	if _, _, err := syscall.Syscall(sysIoctl, fd, ioctlReadTermios, uintptr(unsafe.Pointer(&oldState))); err != 0 {
+		return fmt.Errorf("error when puts the terminal connected to the given file descriptor: %w", err)
 	}
 
 	newState := oldState
 	newState.Lflag &^= syscall.ECHO
 	newState.Lflag |= syscall.ICANON | syscall.ISIG
 	newState.Iflag |= syscall.ICRNL
-	if _, _, e := syscall.Syscall6(sysIoctl, fd, ioctlWriteTermios, uintptr(unsafe.Pointer(&newState)), 0, 0, 0); e != 0 {
-		err = fmt.Errorf("Can't set terminal settings: %v", e)
-		return
+	if _, _, e := syscall.Syscall(sysIoctl, fd, ioctlWriteTermios, uintptr(unsafe.Pointer(&newState))); e != 0 {
+		return fmt.Errorf("error update terminal settings: %w", e)
 	}
-	return
+	return nil
 }
 
-func unlockEcho() (err error) {
+func unlockEcho() error {
 	fd := tty.Fd()
-	if _, _, e := syscall.Syscall6(sysIoctl, fd, ioctlWriteTermios, uintptr(unsafe.Pointer(&oldState)), 0, 0, 0); e != 0 {
-		err = fmt.Errorf("Can't set terminal settings")
+	if _, _, err := syscall.Syscall(sysIoctl, fd, ioctlWriteTermios, uintptr(unsafe.Pointer(&oldState))); err != 0 {
+		return fmt.Errorf("error restores the terminal connected to the given file descriptor: %w", err)
 	}
-	return
+	return nil
 }