Add simple byte-counting export progress-bar

Author: ribasushi

core/commands/dag/dag.go

@@ -1,10 +1,13 @@
 package dagcmd
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"math"
+	"os"
 	"strings"
+	"time"
 
 	"github.com/ipfs/go-ipfs/core/commands/cmdenv"
 	"github.com/ipfs/go-ipfs/core/coredag"
@@ -23,6 +26,12 @@ import (
 	//gipfree "github.com/ipld/go-ipld-prime/impl/free"
 	//gipselector "github.com/ipld/go-ipld-prime/traversal/selector"
 	//gipselectorbuilder "github.com/ipld/go-ipld-prime/traversal/selector/builder"
+
+	"gopkg.in/cheggaaa/pb.v1"
+)
+
+const (
+	progressOptionName = "progress"
 )
 
 var DagCmd = &cmds.Command{
@@ -261,6 +270,9 @@ The output of blocks happens in strict DAG-traversal, first-seen, order.
 	Arguments: []cmds.Argument{
 		cmds.StringArg("root", true, false, "CID of a root to recursively export").EnableStdin(),
 	},
+	Options: []cmds.Option{
+		cmds.BoolOption(progressOptionName, "p", "Display progress on CLI. Defaults to true when STDERR is a TTY."),
+	},
 	Run: func(req *cmds.Request, res cmds.ResponseEmitter, env cmds.Environment) error {
 
 		c, err := cid.Decode(req.Arguments[0])
@@ -334,4 +346,61 @@ The output of blocks happens in strict DAG-traversal, first-seen, order.
 
 		return err
 	},
+	PostRun: cmds.PostRunMap{
+		cmds.CLI: func(res cmds.Response, re cmds.ResponseEmitter) error {
+
+			var showProgress bool
+			val, specified := res.Request().Options[progressOptionName]
+			if !specified {
+				// default based on TTY availability
+				errStat, _ := os.Stderr.Stat()
+				if 0 != (errStat.Mode() & os.ModeCharDevice) {
+					showProgress = true
+				}
+			} else if val.(bool) {
+				showProgress = true
+			}
+
+			// simple passthrough, no progress
+			if !showProgress {
+				return cmds.Copy(re, res)
+			}
+
+			bar := pb.New64(0).SetUnits(pb.U_BYTES)
+			bar.Output = os.Stderr
+			bar.ShowSpeed = true
+			bar.ShowElapsedTime = true
+			bar.RefreshRate = 500 * time.Millisecond
+			bar.Start()
+
+			var processedOneResponse bool
+			for {
+				v, err := res.Next()
+				if err == io.EOF {
+
+					// We only write the final bar update on success
+					// On error it looks too weird
+					bar.Finish()
+
+					return re.Close()
+				} else if err != nil {
+					return re.CloseWithError(err)
+				} else if processedOneResponse {
+					return re.CloseWithError(errors.New("unexpected multipart response during emit, please file a bugreport"))
+				}
+
+				r, ok := v.(io.Reader)
+				if !ok {
+					// some sort of encoded response, this should not be happening
+					return errors.New("unexpected non-stream passed to PostRun: please file a bugreport")
+				}
+
+				processedOneResponse = true
+
+				if err := re.Emit(bar.NewProxyReader(r)); err != nil {
+					return err
+				}
+			}
+		},
+	},
 }