ollama/progress/spinner.go

74 lines
1.2 KiB
Go
Raw Normal View History

2023-11-15 00:33:16 +00:00
package progress
import (
"fmt"
"strings"
"time"
)
type Spinner struct {
message string
messageWidth int
parts []string
value int
ticker *time.Ticker
started time.Time
stopped time.Time
}
func NewSpinner(message string) *Spinner {
s := &Spinner{
message: message,
parts: []string{
"⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏",
},
started: time.Now(),
}
go s.start()
return s
}
func (s *Spinner) String() string {
var sb strings.Builder
2023-11-15 00:33:16 +00:00
if len(s.message) > 0 {
message := strings.TrimSpace(s.message)
if s.messageWidth > 0 && len(message) > s.messageWidth {
message = message[:s.messageWidth]
}
fmt.Fprintf(&sb, "%s", message)
if s.messageWidth-sb.Len() >= 0 {
sb.WriteString(strings.Repeat(" ", s.messageWidth-sb.Len()))
2023-11-15 00:33:16 +00:00
}
sb.WriteString(" ")
2023-11-15 00:33:16 +00:00
}
if s.stopped.IsZero() {
spinner := s.parts[s.value]
sb.WriteString(spinner)
sb.WriteString(" ")
2023-11-15 00:33:16 +00:00
}
return sb.String()
2023-11-15 00:33:16 +00:00
}
func (s *Spinner) start() {
s.ticker = time.NewTicker(100 * time.Millisecond)
for range s.ticker.C {
s.value = (s.value + 1) % len(s.parts)
if !s.stopped.IsZero() {
return
}
}
}
func (s *Spinner) Stop() {
if s.stopped.IsZero() {
s.stopped = time.Now()
}
}