You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
264 lines
7.3 KiB
264 lines
7.3 KiB
// Copyright 2016 CoreOS Inc
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package progressutil
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
|
|
"golang.org/x/crypto/ssh/terminal"
|
|
)
|
|
|
|
var (
|
|
// ErrorProgressOutOfBounds is returned if the progress is set to a value
|
|
// not between 0 and 1.
|
|
ErrorProgressOutOfBounds = fmt.Errorf("progress is out of bounds (0 to 1)")
|
|
|
|
// ErrorNoBarsAdded is returned when no progress bars have been added to a
|
|
// ProgressBarPrinter before PrintAndWait is called.
|
|
ErrorNoBarsAdded = fmt.Errorf("AddProgressBar hasn't been called yet")
|
|
)
|
|
|
|
// ProgressBar represents one progress bar in a ProgressBarPrinter. Should not
|
|
// be created directly, use the AddProgressBar on a ProgressBarPrinter to
|
|
// create these.
|
|
type ProgressBar struct {
|
|
lock sync.Mutex
|
|
|
|
currentProgress float64
|
|
printBefore string
|
|
printAfter string
|
|
done bool
|
|
}
|
|
|
|
func (pb *ProgressBar) clone() *ProgressBar {
|
|
pb.lock.Lock()
|
|
pbClone := &ProgressBar{
|
|
currentProgress: pb.currentProgress,
|
|
printBefore: pb.printBefore,
|
|
printAfter: pb.printAfter,
|
|
done: pb.done,
|
|
}
|
|
pb.lock.Unlock()
|
|
return pbClone
|
|
}
|
|
|
|
func (pb *ProgressBar) GetCurrentProgress() float64 {
|
|
pb.lock.Lock()
|
|
val := pb.currentProgress
|
|
pb.lock.Unlock()
|
|
return val
|
|
}
|
|
|
|
// SetCurrentProgress sets the progress of this ProgressBar. The progress must
|
|
// be between 0 and 1 inclusive.
|
|
func (pb *ProgressBar) SetCurrentProgress(progress float64) error {
|
|
if progress < 0 || progress > 1 {
|
|
return ErrorProgressOutOfBounds
|
|
}
|
|
pb.lock.Lock()
|
|
pb.currentProgress = progress
|
|
pb.lock.Unlock()
|
|
return nil
|
|
}
|
|
|
|
// GetDone returns whether or not this progress bar is done
|
|
func (pb *ProgressBar) GetDone() bool {
|
|
pb.lock.Lock()
|
|
val := pb.done
|
|
pb.lock.Unlock()
|
|
return val
|
|
}
|
|
|
|
// SetDone sets whether or not this progress bar is done
|
|
func (pb *ProgressBar) SetDone(val bool) {
|
|
pb.lock.Lock()
|
|
pb.done = val
|
|
pb.lock.Unlock()
|
|
}
|
|
|
|
// GetPrintBefore gets the text printed on the line before the progress bar.
|
|
func (pb *ProgressBar) GetPrintBefore() string {
|
|
pb.lock.Lock()
|
|
val := pb.printBefore
|
|
pb.lock.Unlock()
|
|
return val
|
|
}
|
|
|
|
// SetPrintBefore sets the text printed on the line before the progress bar.
|
|
func (pb *ProgressBar) SetPrintBefore(before string) {
|
|
pb.lock.Lock()
|
|
pb.printBefore = before
|
|
pb.lock.Unlock()
|
|
}
|
|
|
|
// GetPrintAfter gets the text printed on the line after the progress bar.
|
|
func (pb *ProgressBar) GetPrintAfter() string {
|
|
pb.lock.Lock()
|
|
val := pb.printAfter
|
|
pb.lock.Unlock()
|
|
return val
|
|
}
|
|
|
|
// SetPrintAfter sets the text printed on the line after the progress bar.
|
|
func (pb *ProgressBar) SetPrintAfter(after string) {
|
|
pb.lock.Lock()
|
|
pb.printAfter = after
|
|
pb.lock.Unlock()
|
|
}
|
|
|
|
// ProgressBarPrinter will print out the progress of some number of
|
|
// ProgressBars.
|
|
type ProgressBarPrinter struct {
|
|
lock sync.Mutex
|
|
|
|
// DisplayWidth can be set to influence how large the progress bars are.
|
|
// The bars will be scaled to attempt to produce lines of this number of
|
|
// characters, but lines of different lengths may still be printed. When
|
|
// this value is 0 (aka unset), 80 character columns are assumed.
|
|
DisplayWidth int
|
|
// PadToBeEven, when set to true, will make Print pad the printBefore text
|
|
// with trailing spaces and the printAfter text with leading spaces to make
|
|
// the progress bars the same length.
|
|
PadToBeEven bool
|
|
numLinesInLastPrint int
|
|
progressBars []*ProgressBar
|
|
maxBefore int
|
|
maxAfter int
|
|
|
|
// printToTTYAlways forces this ProgressBarPrinter to always behave as if
|
|
// in a tty. Used for tests.
|
|
printToTTYAlways bool
|
|
}
|
|
|
|
// AddProgressBar will create a new ProgressBar, register it with this
|
|
// ProgressBarPrinter, and return it. This must be called at least once before
|
|
// PrintAndWait is called.
|
|
func (pbp *ProgressBarPrinter) AddProgressBar() *ProgressBar {
|
|
pb := &ProgressBar{}
|
|
pbp.lock.Lock()
|
|
pbp.progressBars = append(pbp.progressBars, pb)
|
|
pbp.lock.Unlock()
|
|
return pb
|
|
}
|
|
|
|
// Print will print out progress information for each ProgressBar that has been
|
|
// added to this ProgressBarPrinter. The progress will be written to printTo,
|
|
// and if printTo is a terminal it will draw progress bars. AddProgressBar
|
|
// must be called at least once before Print is called. If printing to a
|
|
// terminal, all draws after the first one will move the cursor up to draw over
|
|
// the previously printed bars.
|
|
func (pbp *ProgressBarPrinter) Print(printTo io.Writer) (bool, error) {
|
|
pbp.lock.Lock()
|
|
var bars []*ProgressBar
|
|
for _, bar := range pbp.progressBars {
|
|
bars = append(bars, bar.clone())
|
|
}
|
|
numColumns := pbp.DisplayWidth
|
|
pbp.lock.Unlock()
|
|
|
|
if len(bars) == 0 {
|
|
return false, ErrorNoBarsAdded
|
|
}
|
|
|
|
if numColumns == 0 {
|
|
numColumns = 80
|
|
}
|
|
|
|
if pbp.isTerminal(printTo) {
|
|
moveCursorUp(printTo, pbp.numLinesInLastPrint)
|
|
}
|
|
|
|
for _, bar := range bars {
|
|
beforeSize := len(bar.GetPrintBefore())
|
|
afterSize := len(bar.GetPrintAfter())
|
|
if beforeSize > pbp.maxBefore {
|
|
pbp.maxBefore = beforeSize
|
|
}
|
|
if afterSize > pbp.maxAfter {
|
|
pbp.maxAfter = afterSize
|
|
}
|
|
}
|
|
|
|
allDone := true
|
|
for _, bar := range bars {
|
|
if pbp.isTerminal(printTo) {
|
|
bar.printToTerminal(printTo, numColumns, pbp.PadToBeEven, pbp.maxBefore, pbp.maxAfter)
|
|
} else {
|
|
bar.printToNonTerminal(printTo)
|
|
}
|
|
allDone = allDone && bar.GetCurrentProgress() == 1
|
|
}
|
|
|
|
pbp.numLinesInLastPrint = len(bars)
|
|
|
|
return allDone, nil
|
|
}
|
|
|
|
// moveCursorUp moves the cursor up numLines in the terminal
|
|
func moveCursorUp(printTo io.Writer, numLines int) {
|
|
if numLines > 0 {
|
|
fmt.Fprintf(printTo, "\033[%dA", numLines)
|
|
}
|
|
}
|
|
|
|
func (pb *ProgressBar) printToTerminal(printTo io.Writer, numColumns int, padding bool, maxBefore, maxAfter int) {
|
|
before := pb.GetPrintBefore()
|
|
after := pb.GetPrintAfter()
|
|
|
|
if padding {
|
|
before = before + strings.Repeat(" ", maxBefore-len(before))
|
|
after = strings.Repeat(" ", maxAfter-len(after)) + after
|
|
}
|
|
|
|
progressBarSize := numColumns - (len(fmt.Sprintf("%s [] %s", before, after)))
|
|
progressBar := ""
|
|
if progressBarSize > 0 {
|
|
currentProgress := int(pb.GetCurrentProgress() * float64(progressBarSize))
|
|
progressBar = fmt.Sprintf("[%s%s] ",
|
|
strings.Repeat("=", currentProgress),
|
|
strings.Repeat(" ", progressBarSize-currentProgress))
|
|
} else {
|
|
// If we can't fit the progress bar, better to not pad the before/after.
|
|
before = pb.GetPrintBefore()
|
|
after = pb.GetPrintAfter()
|
|
}
|
|
|
|
fmt.Fprintf(printTo, "%s %s%s\n", before, progressBar, after)
|
|
}
|
|
|
|
func (pb *ProgressBar) printToNonTerminal(printTo io.Writer) {
|
|
if !pb.GetDone() {
|
|
fmt.Fprintf(printTo, "%s %s\n", pb.printBefore, pb.printAfter)
|
|
if pb.GetCurrentProgress() == 1 {
|
|
pb.SetDone(true)
|
|
}
|
|
}
|
|
}
|
|
|
|
// isTerminal returns True when w is going to a tty, and false otherwise.
|
|
func (pbp *ProgressBarPrinter) isTerminal(w io.Writer) bool {
|
|
if pbp.printToTTYAlways {
|
|
return true
|
|
}
|
|
if f, ok := w.(*os.File); ok {
|
|
return terminal.IsTerminal(int(f.Fd()))
|
|
}
|
|
return false
|
|
}
|