Merge pull request #151 from nilsbeck/feature/allow-custom-spinners

feature/allow-custom-spinners

Author: schollz

examples/customization/main.go

@@ -38,5 +38,5 @@ func main() {
 
 	// got notified that progress bar is complete.
 	<-doneCh
-	fmt.Println("\n ======= progress bar completed ==========\n")
+	fmt.Println("\n ======= progress bar completed ==========")
 }

examples/download/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"io"
 	"net/http"
 	"os"
@@ -11,7 +12,7 @@ import (
 func main() {
 	req, _ := http.NewRequest("GET", "https://dl.google.com/go/go1.14.2.src.tar.gz", nil)
 	resp, _ := http.DefaultClient.Do(req)
-	defer resp.Body.Close()
+	defer check(resp.Body.Close)
 
 	f, _ := os.OpenFile("go1.14.2.src.tar.gz", os.O_CREATE|os.O_WRONLY, 0644)
 	defer f.Close()
@@ -22,3 +23,10 @@ func main() {
 	)
 	io.Copy(io.MultiWriter(f, bar), resp.Body)
 }
+
+// check checks the returned error of a function.
+func check(f func() error) {
+	if err := f(); err != nil {
+		fmt.Fprintf(os.Stderr, "received error: %v\n", err)
+	}
+}

examples/pacman/main.go

@@ -37,5 +37,5 @@ func main() {
 
 	// got notified that progress bar is complete.
 	<-doneCh
-	fmt.Println("\n ======= progress bar completed ==========\n")
+	fmt.Println("\n ======= progress bar completed ==========")
 }

progressbar.go

@@ -96,6 +96,12 @@ type config struct {
 	// spinnerType should be a number between 0-75
 	spinnerType int
 
+	// spinnerTypeOptionUsed remembers if the spinnerType was changed manually
+	spinnerTypeOptionUsed bool
+
+	// spinner represents the spinner as a slice of string
+	spinner []string
+
 	// fullWidth specifies whether to measure and set the bar to a specific width
 	fullWidth bool
 
@@ -134,10 +140,19 @@ func OptionSetWidth(s int) Option {
 // OptionSpinnerType sets the type of spinner used for indeterminate bars
 func OptionSpinnerType(spinnerType int) Option {
 	return func(p *ProgressBar) {
+		p.config.spinnerTypeOptionUsed = true
 		p.config.spinnerType = spinnerType
 	}
 }
 
+// OptionSpinnerCustom sets the spinner used for indeterminate bars to the passed
+// slice of string
+func OptionSpinnerCustom(spinner []string) Option {
+	return func(p *ProgressBar) {
+		p.config.spinner = spinner
+	}
+}
+
 // OptionSetTheme sets the elements the bar is constructed of
 func OptionSetTheme(t Theme) Option {
 	return func(p *ProgressBar) {
@@ -509,6 +524,11 @@ func (p *ProgressBar) Add64(num int64) error {
 		return nil
 	}
 
+	// error out since OptionSpinnerCustom will always override a manually set spinnerType
+	if p.config.spinnerTypeOptionUsed && len(p.config.spinner) > 0 {
+		return errors.New("OptionSpinnerType and OptionSpinnerCustom cannot be used together")
+	}
+
 	if p.config.max == 0 {
 		return errors.New("max must be greater than 0")
 	}
@@ -853,7 +873,11 @@ func renderProgressBar(c config, s *state) (int, error) {
 	str := ""
 
 	if c.ignoreLength {
-		spinner := spinners[c.spinnerType][int(math.Round(math.Mod(float64(time.Since(s.startTime).Milliseconds()/100), float64(len(spinners[c.spinnerType])))))]
+		selectedSpinner := spinners[c.spinnerType]
+		if len(c.spinner) > 0 {
+			selectedSpinner = c.spinner
+		}
+		spinner := selectedSpinner[int(math.Round(math.Mod(float64(time.Since(s.startTime).Milliseconds()/100), float64(len(selectedSpinner)))))]
 		if c.elapsedTime {
 			if c.showDescriptionAtLineEnd {
 				str = fmt.Sprintf("\r%s %s [%s] %s ",

progressbar_test.go

@@ -201,13 +201,56 @@ func TestSpinnerType(t *testing.T) {
 	bar.Reset()
 	for i := 0; i < 10; i++ {
 		time.Sleep(120 * time.Millisecond)
-		bar.Add(1)
+		err := bar.Add(1)
+		if err != nil {
+			t.Errorf("Successfully tested one spinner option can be used.")
+		}
 	}
 	if false {
 		t.Errorf("error")
 	}
 }
 
+func TestSpinnerCustom(t *testing.T) {
+	bar := NewOptions(-1,
+		OptionSetWidth(10),
+		OptionSetDescription("indeterminate spinner"),
+		OptionShowIts(),
+		OptionShowCount(),
+		OptionSpinnerCustom([]string{"🐰", "🐰", "🥕", "🥕"}),
+	)
+	bar.Reset()
+	for i := 0; i < 10; i++ {
+		time.Sleep(120 * time.Millisecond)
+		err := bar.Add(1)
+		if err != nil {
+			t.Errorf("Successfully tested one spinner option can be used.")
+		}
+	}
+	if false {
+		t.Errorf("error")
+	}
+}
+
+func TestSpinnerTypeAndCustom(t *testing.T) {
+	bar := NewOptions(-1,
+		OptionSetWidth(10),
+		OptionSetDescription("indeterminate spinner"),
+		OptionShowIts(),
+		OptionShowCount(),
+		OptionSpinnerCustom([]string{"🐰", "🐰", "🥕", "🥕"}),
+		OptionSpinnerType(9),
+	)
+	bar.Reset()
+	for i := 0; i < 10; i++ {
+		time.Sleep(120 * time.Millisecond)
+		err := bar.Add(1)
+		if err == nil {
+			t.Errorf("Successfully tested both spinner options cannot be used together.")
+		}
+	}
+}
+
 func Test_IsFinished(t *testing.T) {
 	isCalled := false
 	bar := NewOptions(72, OptionOnCompletion(func() {