package console import ( "fmt" "github.com/cheggaaa/pb" "github.com/smira/aptly/aptly" "github.com/wsxiaoys/terminal/color" "strings" ) const ( codePrint = iota codeProgress codeHideProgress codeStop codeFlush ) type printTask struct { code int message string reply chan bool } // Progress is a progress displaying subroutine, it allows to show download and other operations progress // mixed with progress bar type Progress struct { stop chan bool stopped chan bool queue chan printTask bar *pb.ProgressBar barShown bool } // Check interface var ( _ aptly.Progress = (*Progress)(nil) ) // NewProgress creates new progress instance func NewProgress() *Progress { return &Progress{ stopped: make(chan bool), queue: make(chan printTask, 100), } } // Start makes progress start its work func (p *Progress) Start() { go p.worker() } // Shutdown shuts down progress display func (p *Progress) Shutdown() { p.ShutdownBar() p.queue <- printTask{code: codeStop} <-p.stopped } // Flush waits for all queued messages to be displayed func (p *Progress) Flush() { ch := make(chan bool) p.queue <- printTask{code: codeFlush, reply: ch} <-ch } // InitBar starts progressbar for count bytes or count items func (p *Progress) InitBar(count int64, isBytes bool) { if p.bar != nil { panic("bar already initialized") } if RunningOnTerminal() { p.bar = pb.New(0) p.bar.Total = count p.bar.NotPrint = true p.bar.Callback = func(out string) { p.queue <- printTask{code: codeProgress, message: out} } if isBytes { p.bar.SetUnits(pb.U_BYTES) p.bar.ShowSpeed = true } p.bar.Start() } } // ShutdownBar stops progress bar and hides it func (p *Progress) ShutdownBar() { if p.bar == nil { return } p.bar.Finish() p.bar = nil p.queue <- printTask{code: codeHideProgress} } // Write is implementation of io.Writer to support updating of progress bar func (p *Progress) Write(s []byte) (int, error) { if p.bar != nil { p.bar.Add(len(s)) } return len(s), nil } // AddBar increments progress for progress bar func (p *Progress) AddBar(count int) { if p.bar != nil { p.bar.Add(count) } } // SetBar sets current position for progress bar func (p *Progress) SetBar(count int) { if p.bar != nil { p.bar.Set(count) } } // Printf does printf but in safe manner: not overwriting progress bar func (p *Progress) Printf(msg string, a ...interface{}) { p.queue <- printTask{code: codePrint, message: fmt.Sprintf(msg, a...)} } // ColoredPrintf does printf in colored way + newline func (p *Progress) ColoredPrintf(msg string, a ...interface{}) { if RunningOnTerminal() { p.queue <- printTask{code: codePrint, message: color.Sprintf(msg, a...) + "\n"} } else { // stip color marks var prev rune msg = strings.Map(func(r rune) rune { if prev == '@' { prev = 0 if r == '@' { return r } return -1 } prev = r if r == '@' { return -1 } return r }, msg) p.Printf(msg+"\n", a...) } } func (p *Progress) worker() { for { task := <-p.queue switch task.code { case codePrint: if p.barShown { fmt.Print("\r\033[2K") p.barShown = false } fmt.Print(task.message) case codeProgress: if p.bar != nil { fmt.Print("\r" + task.message) p.barShown = true } case codeHideProgress: if p.barShown { fmt.Print("\r\033[2K") p.barShown = false } case codeFlush: task.reply <- true case codeStop: p.stopped <- true return } } }