github.com/1aal/kubeblocks@v0.0.0-20231107070852-e1c03e598921/pkg/cli/spinner/spinner.go (about)

     1  /*
     2  Copyright (C) 2022-2023 ApeCloud Co., Ltd
     3  
     4  This file is part of KubeBlocks project
     5  
     6  This program is free software: you can redistribute it and/or modify
     7  it under the terms of the GNU Affero General Public License as published by
     8  the Free Software Foundation, either version 3 of the License, or
     9  (at your option) any later version.
    10  
    11  This program is distributed in the hope that it will be useful
    12  but WITHOUT ANY WARRANTY; without even the implied warranty of
    13  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    14  GNU Affero General Public License for more details.
    15  
    16  You should have received a copy of the GNU Affero General Public License
    17  along with this program.  If not, see <http://www.gnu.org/licenses/>.
    18  */
    19  
    20  package spinner
    21  
    22  import (
    23  	"fmt"
    24  	"io"
    25  	"os"
    26  	"os/signal"
    27  	"strings"
    28  	"syscall"
    29  	"time"
    30  
    31  	"github.com/briandowns/spinner"
    32  
    33  	"github.com/1aal/kubeblocks/pkg/cli/printer"
    34  	"github.com/1aal/kubeblocks/pkg/cli/util"
    35  )
    36  
    37  type Spinner struct {
    38  	s      *spinner.Spinner
    39  	delay  time.Duration
    40  	cancel chan struct{}
    41  }
    42  
    43  type Interface interface {
    44  	Start()
    45  	Done(status string)
    46  	Success()
    47  	Fail()
    48  	SetMessage(msg string)
    49  	SetFinalMsg(msg string)
    50  	updateSpinnerMessage(msg string)
    51  }
    52  
    53  type Option func(Interface)
    54  
    55  func WithMessage(msg string) Option {
    56  	return func(s Interface) {
    57  		s.updateSpinnerMessage(msg)
    58  	}
    59  }
    60  
    61  func (s *Spinner) updateSpinnerMessage(msg string) {
    62  	s.s.Suffix = fmt.Sprintf(" %s", msg)
    63  }
    64  
    65  func (s *Spinner) SetMessage(msg string) {
    66  	s.updateSpinnerMessage(msg)
    67  	if !s.s.Active() {
    68  		s.Start()
    69  	}
    70  }
    71  
    72  func (s *Spinner) Start() {
    73  	if s.cancel != nil {
    74  		return
    75  	}
    76  	if s.delay == 0 {
    77  		s.s.Start()
    78  		return
    79  	}
    80  	s.cancel = make(chan struct{}, 1)
    81  	go func() {
    82  		select {
    83  		case <-s.cancel:
    84  			return
    85  		case <-time.After(s.delay):
    86  			s.s.Start()
    87  			s.cancel = nil
    88  		}
    89  		time.Sleep(50 * time.Millisecond)
    90  	}()
    91  }
    92  
    93  func (s *Spinner) Done(status string) {
    94  	if s.cancel != nil {
    95  		close(s.cancel)
    96  	}
    97  	s.stop(status)
    98  }
    99  
   100  func (s *Spinner) SetFinalMsg(msg string) {
   101  	s.s.FinalMSG = msg
   102  	if !s.s.Active() {
   103  		s.Start()
   104  	}
   105  }
   106  
   107  func (s *Spinner) stop(status string) {
   108  	if s.s == nil {
   109  		return
   110  	}
   111  
   112  	if status != "" {
   113  		s.s.FinalMSG = fmt.Sprintf("%s %s\n", strings.TrimPrefix(s.s.Suffix, " "), status)
   114  	}
   115  	s.s.Stop()
   116  
   117  	// show cursor in terminal.
   118  	fmt.Fprintf(s.s.Writer, "\033[?25h")
   119  }
   120  
   121  func (s *Spinner) Success() {
   122  	s.Done(printer.BoldGreen("OK"))
   123  }
   124  
   125  func (s *Spinner) Fail() {
   126  	s.Done(printer.BoldRed("FAIL"))
   127  }
   128  
   129  func New(w io.Writer, opts ...Option) Interface {
   130  	if util.IsWindows() {
   131  		return NewWindowsSpinner(w, opts...)
   132  	}
   133  
   134  	res := &Spinner{}
   135  	res.s = spinner.New(spinner.CharSets[11],
   136  		100*time.Millisecond,
   137  		spinner.WithWriter(w),
   138  		spinner.WithHiddenCursor(true),
   139  		spinner.WithColor("cyan"),
   140  	)
   141  
   142  	for _, opt := range opts {
   143  		opt(res)
   144  	}
   145  
   146  	c := make(chan os.Signal, 1)
   147  	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
   148  	// Capture the interrupt signal, make the `spinner` program exit gracefully, and prevent the cursor from disappearing.
   149  	go func() {
   150  		<-c
   151  		res.Done("")
   152  		os.Exit(0)
   153  	}()
   154  	res.Start()
   155  	return res
   156  }