blob: 7b7b081dbe02da54c546c93e5967043d7f1c7be9 [file]
// Copyright (c) 2021 Brian J. Downs
//
// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner
import (
"errors"
"fmt"
"io"
"math"
"os"
"runtime"
"strconv"
"strings"
"sync"
"time"
"unicode/utf8"
"github.com/fatih/color"
"github.com/mattn/go-isatty"
"golang.org/x/term"
)
// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")
// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
// default colors for backwards compatibility
"black": true,
"red": true,
"green": true,
"yellow": true,
"blue": true,
"magenta": true,
"cyan": true,
"white": true,
// attributes
"reset": true,
"bold": true,
"faint": true,
"italic": true,
"underline": true,
"blinkslow": true,
"blinkrapid": true,
"reversevideo": true,
"concealed": true,
"crossedout": true,
// foreground text
"fgBlack": true,
"fgRed": true,
"fgGreen": true,
"fgYellow": true,
"fgBlue": true,
"fgMagenta": true,
"fgCyan": true,
"fgWhite": true,
// foreground Hi-Intensity text
"fgHiBlack": true,
"fgHiRed": true,
"fgHiGreen": true,
"fgHiYellow": true,
"fgHiBlue": true,
"fgHiMagenta": true,
"fgHiCyan": true,
"fgHiWhite": true,
// background text
"bgBlack": true,
"bgRed": true,
"bgGreen": true,
"bgYellow": true,
"bgBlue": true,
"bgMagenta": true,
"bgCyan": true,
"bgWhite": true,
// background Hi-Intensity text
"bgHiBlack": true,
"bgHiRed": true,
"bgHiGreen": true,
"bgHiYellow": true,
"bgHiBlue": true,
"bgHiMagenta": true,
"bgHiCyan": true,
"bgHiWhite": true,
}
// returns true if the OS is windows and the WT_SESSION env variable is set.
var isWindows = runtime.GOOS == "windows"
var isWindowsTerminalOnWindows = len(os.Getenv("WT_SESSION")) > 0 && isWindows
// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
// default colors for backwards compatibility
"black": color.FgBlack,
"red": color.FgRed,
"green": color.FgGreen,
"yellow": color.FgYellow,
"blue": color.FgBlue,
"magenta": color.FgMagenta,
"cyan": color.FgCyan,
"white": color.FgWhite,
// attributes
"reset": color.Reset,
"bold": color.Bold,
"faint": color.Faint,
"italic": color.Italic,
"underline": color.Underline,
"blinkslow": color.BlinkSlow,
"blinkrapid": color.BlinkRapid,
"reversevideo": color.ReverseVideo,
"concealed": color.Concealed,
"crossedout": color.CrossedOut,
// foreground text colors
"fgBlack": color.FgBlack,
"fgRed": color.FgRed,
"fgGreen": color.FgGreen,
"fgYellow": color.FgYellow,
"fgBlue": color.FgBlue,
"fgMagenta": color.FgMagenta,
"fgCyan": color.FgCyan,
"fgWhite": color.FgWhite,
// foreground Hi-Intensity text colors
"fgHiBlack": color.FgHiBlack,
"fgHiRed": color.FgHiRed,
"fgHiGreen": color.FgHiGreen,
"fgHiYellow": color.FgHiYellow,
"fgHiBlue": color.FgHiBlue,
"fgHiMagenta": color.FgHiMagenta,
"fgHiCyan": color.FgHiCyan,
"fgHiWhite": color.FgHiWhite,
// background text colors
"bgBlack": color.BgBlack,
"bgRed": color.BgRed,
"bgGreen": color.BgGreen,
"bgYellow": color.BgYellow,
"bgBlue": color.BgBlue,
"bgMagenta": color.BgMagenta,
"bgCyan": color.BgCyan,
"bgWhite": color.BgWhite,
// background Hi-Intensity text colors
"bgHiBlack": color.BgHiBlack,
"bgHiRed": color.BgHiRed,
"bgHiGreen": color.BgHiGreen,
"bgHiYellow": color.BgHiYellow,
"bgHiBlue": color.BgHiBlue,
"bgHiMagenta": color.BgHiMagenta,
"bgHiCyan": color.BgHiCyan,
"bgHiWhite": color.BgHiWhite,
}
// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
return validColors[c]
}
// Spinner struct to hold the provided options.
type Spinner struct {
mu *sync.RWMutex
Delay time.Duration // Delay is the speed of the indicator
chars []string // chars holds the chosen character set
Prefix string // Prefix is the text preppended to the indicator
Suffix string // Suffix is the text appended to the indicator
FinalMSG string // string displayed after Stop() is called
lastOutputPlain string // last character(set) written
LastOutput string // last character(set) written with colors
color func(a ...interface{}) string // default color is white
Writer io.Writer // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
active bool // active holds the state of the spinner
enabled bool // indicates whether the spinner is enabled or not
stopChan chan struct{} // stopChan is a channel used to stop the indicator
HideCursor bool // hideCursor determines if the cursor is visible
PreUpdate func(s *Spinner) // will be triggered before every spinner update
PostUpdate func(s *Spinner) // will be triggered after every spinner update
}
// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
s := &Spinner{
Delay: d,
chars: cs,
color: color.New(color.FgWhite).SprintFunc(),
mu: &sync.RWMutex{},
Writer: color.Output,
stopChan: make(chan struct{}, 1),
active: false,
enabled: true,
HideCursor: true,
}
for _, option := range options {
option(s)
}
return s
}
// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)
// Options contains fields to configure the spinner.
type Options struct {
Color string
Suffix string
FinalMSG string
HideCursor bool
}
// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
return func(s *Spinner) {
s.Color(color)
}
}
// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
return func(s *Spinner) {
s.Suffix = suffix
}
}
// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
return func(s *Spinner) {
s.FinalMSG = finalMsg
}
}
// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
return func(s *Spinner) {
s.HideCursor = hideCursor
}
}
// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
return func(s *Spinner) {
s.mu.Lock()
s.Writer = w
s.mu.Unlock()
}
}
// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
return s.active
}
// Enabled returns whether or not the spinner is enabled.
func (s *Spinner) Enabled() bool {
return s.enabled
}
// Enable enables and restarts the spinner
func (s *Spinner) Enable() {
s.enabled = true
s.Restart()
}
// Disable stops and disables the spinner
func (s *Spinner) Disable() {
s.enabled = false
s.Stop()
}
// Start will start the indicator.
func (s *Spinner) Start() {
s.mu.Lock()
if s.active || !s.enabled || !isRunningInTerminal() {
s.mu.Unlock()
return
}
if s.HideCursor && !isWindowsTerminalOnWindows {
// hides the cursor
fmt.Fprint(s.Writer, "\033[?25l")
}
// Disable colors for simple Windows CMD or Powershell
// as they can not recognize them
if isWindows && !isWindowsTerminalOnWindows {
color.NoColor = true
}
s.active = true
s.mu.Unlock()
go func() {
for {
for i := 0; i < len(s.chars); i++ {
select {
case <-s.stopChan:
return
default:
s.mu.Lock()
if !s.active {
s.mu.Unlock()
return
}
if !isWindowsTerminalOnWindows {
s.erase()
}
if s.PreUpdate != nil {
s.PreUpdate(s)
}
var outColor string
if isWindows {
if s.Writer == os.Stderr {
outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
} else {
outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
}
} else {
outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
}
outPlain := fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
fmt.Fprint(s.Writer, outColor)
s.lastOutputPlain = outPlain
s.LastOutput = outColor
delay := s.Delay
if s.PostUpdate != nil {
s.PostUpdate(s)
}
s.mu.Unlock()
time.Sleep(delay)
}
}
}
}()
}
// Stop stops the indicator.
func (s *Spinner) Stop() {
s.mu.Lock()
defer s.mu.Unlock()
if s.active {
s.active = false
if s.HideCursor && !isWindowsTerminalOnWindows {
// makes the cursor visible
fmt.Fprint(s.Writer, "\033[?25h")
}
s.erase()
if s.FinalMSG != "" {
if isWindowsTerminalOnWindows {
fmt.Fprint(s.Writer, "\r", s.FinalMSG)
} else {
fmt.Fprint(s.Writer, s.FinalMSG)
}
}
s.stopChan <- struct{}{}
}
}
// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
s.Stop()
s.Start()
}
// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
s.mu.Lock()
for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
}
s.mu.Unlock()
}
// Color will set the struct field for the given color to be used. The spinner
// will need to be explicitly restarted.
func (s *Spinner) Color(colors ...string) error {
colorAttributes := make([]color.Attribute, len(colors))
// Verify colours are valid and place the appropriate attribute in the array
for index, c := range colors {
if !validColor(c) {
return errInvalidColor
}
colorAttributes[index] = colorAttributeMap[c]
}
s.mu.Lock()
s.color = color.New(colorAttributes...).SprintFunc()
s.mu.Unlock()
return nil
}
// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
s.mu.Lock()
s.Delay = d
s.mu.Unlock()
}
// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
s.mu.Lock()
s.chars = cs
s.mu.Unlock()
}
// erase deletes written characters on the current line.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
n := utf8.RuneCountInString(s.lastOutputPlain)
if runtime.GOOS == "windows" && !isWindowsTerminalOnWindows {
clearString := "\r" + strings.Repeat(" ", n) + "\r"
fmt.Fprint(s.Writer, clearString)
s.lastOutputPlain = ""
return
}
numberOfLinesToErase := computeNumberOfLinesNeededToPrintString(s.lastOutputPlain)
// Taken from https://en.wikipedia.org/wiki/ANSI_escape_code:
// \r - Carriage return - Moves the cursor to column zero
// \033[K - Erases part of the line. If n is 0 (or missing), clear from
// cursor to the end of the line. If n is 1, clear from cursor to beginning
// of the line. If n is 2, clear entire line. Cursor position does not
// change.
// \033[F - Go to the beginning of previous line
eraseCodeString := strings.Builder{}
// current position is at the end of the last printed line. Start by erasing current line
eraseCodeString.WriteString("\r\033[K") // start by erasing current line
for i := 1; i < numberOfLinesToErase; i++ {
// For each additional lines, go up one line and erase it.
eraseCodeString.WriteString("\033[F\033[K")
}
fmt.Fprintf(s.Writer, eraseCodeString.String())
s.lastOutputPlain = ""
}
// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
s.mu.Lock()
}
// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
s.mu.Unlock()
}
// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
numSeq := make([]string, length)
for i := 0; i < length; i++ {
numSeq[i] = strconv.Itoa(i)
}
return numSeq
}
// isRunningInTerminal check if stdout file descriptor is terminal
func isRunningInTerminal() bool {
return isatty.IsTerminal(os.Stdout.Fd())
}
func computeNumberOfLinesNeededToPrintString(linePrinted string) int {
terminalWidth := math.MaxInt // assume infinity by default to keep behaviour consistent with what we had before
if term.IsTerminal(0) {
if width, _, err := term.GetSize(0); err == nil {
terminalWidth = width
}
}
return computeNumberOfLinesNeededToPrintStringInternal(linePrinted, terminalWidth)
}
// isAnsiMarker returns if a rune denotes the start of an ANSI sequence
func isAnsiMarker(r rune) bool {
return r == '\x1b'
}
// isAnsiTerminator returns if a rune denotes the end of an ANSI sequence
func isAnsiTerminator(r rune) bool {
return (r >= 0x40 && r <= 0x5a) || (r == 0x5e) || (r >= 0x60 && r <= 0x7e)
}
// computeLineWidth returns the displayed width of a line
func computeLineWidth(line string) int {
width := 0
ansi := false
for _, r := range []rune(line) {
// increase width only when outside of ANSI escape sequences
if ansi || isAnsiMarker(r) {
ansi = !isAnsiTerminator(r)
} else {
width += utf8.RuneLen(r)
}
}
return width
}
func computeNumberOfLinesNeededToPrintStringInternal(linePrinted string, maxLineWidth int) int {
lineCount := 0
for _, line := range strings.Split(linePrinted, "\n") {
lineCount += 1
lineWidth := computeLineWidth(line)
if lineWidth > maxLineWidth {
lineCount += int(float64(lineWidth) / float64(maxLineWidth))
}
}
return lineCount
}