ollama/progress/progress.go
2023-11-17 10:06:19 -08:00

92 lines
1.4 KiB
Go

package progress
import (
"fmt"
"io"
"os"
"strings"
"sync"
"time"
"golang.org/x/term"
)
type State interface {
String() string
}
type Progress struct {
mu sync.Mutex
pos int
w io.Writer
ticker *time.Ticker
states []State
}
func NewProgress(w io.Writer) *Progress {
p := &Progress{pos: -1, w: w}
go p.start()
return p
}
func (p *Progress) Stop() bool {
if p.ticker != nil {
p.ticker.Stop()
p.ticker = nil
p.render()
return true
}
return false
}
func (p *Progress) StopAndClear() bool {
stopped := p.Stop()
if stopped {
termWidth, _, err := term.GetSize(int(os.Stderr.Fd()))
if err != nil {
panic(err)
}
// clear the progress bar by:
// 1. reset to beginning of line
// 2. move up to the first line of the progress bar
// 3. fill the terminal width with spaces
// 4. reset to beginning of line
fmt.Fprintf(p.w, "\r\033[%dA%s\r", p.pos, strings.Repeat(" ", termWidth))
}
return stopped
}
func (p *Progress) Add(key string, state State) {
p.mu.Lock()
defer p.mu.Unlock()
p.states = append(p.states, state)
}
func (p *Progress) render() error {
p.mu.Lock()
defer p.mu.Unlock()
fmt.Fprintf(p.w, "\033[%dA", p.pos)
for _, state := range p.states {
fmt.Fprintln(p.w, state.String())
}
if len(p.states) > 0 {
p.pos = len(p.states)
}
return nil
}
func (p *Progress) start() {
p.ticker = time.NewTicker(100 * time.Millisecond)
for range p.ticker.C {
p.render()
}
}