feat: support disable the progress bar (#190)

Signed-off-by: chlins <[email protected]>

Author: chlins

cmd/root.go

@@ -24,11 +24,12 @@ import (
 	"os/signal"
 	"syscall"
 
-	"github.com/CloudNativeAI/modctl/cmd/modelfile"
-	"github.com/CloudNativeAI/modctl/pkg/config"
-
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+
+	"github.com/CloudNativeAI/modctl/cmd/modelfile"
+	internalpb "github.com/CloudNativeAI/modctl/internal/pb"
+	"github.com/CloudNativeAI/modctl/pkg/config"
 )
 
 var rootConfig *config.Root
@@ -51,6 +52,9 @@ var rootCmd = &cobra.Command{
 				}
 			}()
 		}
+
+		// TODO: need refactor as currently use a global flag to control the progress bar render.
+		internalpb.SetDisableProgress(rootConfig.DisableProgress)
 		return nil
 	},
 }
@@ -83,6 +87,7 @@ func init() {
 	flags.StringVar(&rootConfig.StoargeDir, "storage-dir", rootConfig.StoargeDir, "specify the storage directory for modctl")
 	flags.BoolVar(&rootConfig.Pprof, "pprof", rootConfig.Pprof, "enable pprof")
 	flags.StringVar(&rootConfig.PprofAddr, "pprof-addr", rootConfig.PprofAddr, "specify the address for pprof")
+	flags.BoolVar(&rootConfig.DisableProgress, "no-progress", rootConfig.DisableProgress, "disable progress bar")
 
 	// Bind common flags.
 	if err := viper.BindPFlags(flags); err != nil {

internal/pb/pb.go

@@ -28,6 +28,16 @@ import (
 	"github.com/vbauerster/mpb/v8/decor"
 )
 
+var (
+	// disableProgress is the flag to disable progress bar.
+	disableProgress bool
+)
+
+// SetDisableProgress disables the progress bar.
+func SetDisableProgress(disable bool) {
+	disableProgress = disable
+}
+
 // NormalizePrompt normalizes the prompt string.
 func NormalizePrompt(prompt string) string {
 	return fmt.Sprintf("%s =>", prompt)
@@ -71,6 +81,11 @@ func NewProgressBar(writers ...io.Writer) *ProgressBar {
 
 // Add adds a new progress bar.
 func (p *ProgressBar) Add(prompt, name string, size int64, reader io.Reader) io.Reader {
+	// Return the reader directly if progress is disabled.
+	if disableProgress {
+		return reader
+	}
+
 	p.mu.RLock()
 	oldBar := p.bars[name]
 	p.mu.RUnlock()

pkg/backend/pull.go

@@ -59,6 +59,11 @@ func (b *backend) Pull(ctx context.Context, target string, cfg *config.Pull) err
 		return fmt.Errorf("failed to decode the manifest: %w", err)
 	}
 
+	// TODO: need refactor as currently use a global flag to control the progress bar render.
+	if cfg.DisableProgress {
+		internalpb.SetDisableProgress(true)
+	}
+
 	// create the progress bar to track the progress of push.
 	pb := internalpb.NewProgressBar(cfg.ProgressWriter)
 	pb.Start()

pkg/config/pull.go

@@ -38,6 +38,7 @@ type Pull struct {
 	ExtractFromRemote bool
 	Hooks             PullHooks
 	ProgressWriter    io.Writer
+	DisableProgress   bool
 }
 
 func NewPull() *Pull {
@@ -50,6 +51,7 @@ func NewPull() *Pull {
 		ExtractFromRemote: false,
 		Hooks:             &emptyPullHook{},
 		ProgressWriter:    os.Stdout,
+		DisableProgress:   false,
 	}
 }
 

pkg/config/root.go

@@ -22,9 +22,10 @@ import (
 )
 
 type Root struct {
-	StoargeDir string
-	Pprof      bool
-	PprofAddr  string
+	StoargeDir      string
+	Pprof           bool
+	PprofAddr       string
+	DisableProgress bool
 }
 
 func NewRoot() (*Root, error) {
@@ -34,8 +35,9 @@ func NewRoot() (*Root, error) {
 	}
 
 	return &Root{
-		StoargeDir: filepath.Join(user.HomeDir, ".modctl"),
-		Pprof:      false,
-		PprofAddr:  "localhost:6060",
+		StoargeDir:      filepath.Join(user.HomeDir, ".modctl"),
+		Pprof:           false,
+		PprofAddr:       "localhost:6060",
+		DisableProgress: false,
 	}, nil
 }