ls: rework ncdu output to use walker.LeaveDir

cmd/restic/cmd_ls.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 	"io"
 	"os"
-	"path/filepath"
 	"strings"
 	"time"
 
@@ -67,7 +66,7 @@ func init() {
 	flags.BoolVarP(&lsOptions.ListLong, "long", "l", false, "use a long listing format showing size and mode")
 	flags.BoolVar(&lsOptions.Recursive, "recursive", false, "include files in subfolders of the listed directories")
 	flags.BoolVar(&lsOptions.HumanReadable, "human-readable", false, "print sizes in human readable format")
-	flags.BoolVar(&lsOptions.Ncdu, "ncdu", false, "output NCDU save format (pipe into ncdu -f - ")
+	flags.BoolVar(&lsOptions.Ncdu, "ncdu", false, "output NCDU save format (pipe into 'ncdu -f -')")
 }
 
 type lsSnapshot struct {
@@ -119,10 +118,15 @@ func lsNodeJSON(enc *json.Encoder, path string, node *restic.Node) error {
 	return enc.Encode(n)
 }
 
+type ncduPrinter struct {
+	out   io.Writer
+	depth int
+}
+
 // lsSnapshotNcdu prints a restic snapshot in Ncdu save format.
 // It opens the JSON list. Nodes are added with lsNodeNcdu and the list is closed by lsCloseNcdu.
 // Format documentation: https://dev.yorhel.nl/ncdu/jsonfmt
-func lsSnapshotNcdu(stdout io.Writer, depth *int, sn *restic.Snapshot) {
+func (p *ncduPrinter) ProcessSnapshot(sn *restic.Snapshot) {
 	const NcduMajorVer = 1
 	const NcduMinorVer = 2
 
@@ -130,11 +134,11 @@ func lsSnapshotNcdu(stdout io.Writer, depth *int, sn *restic.Snapshot) {
 	if err != nil {
 		Warnf("JSON encode failed: %v\n", err)
 	}
-	*depth++
-	fmt.Fprintf(stdout, "[%d, %d, %s", NcduMajorVer, NcduMinorVer, string(snapshotBytes))
+	p.depth++
+	fmt.Fprintf(p.out, "[%d, %d, %s", NcduMajorVer, NcduMinorVer, string(snapshotBytes))
 }
 
-func lsNodeNcdu(stdout io.Writer, depth *int, currentPath *string, path string, node *restic.Node) {
+func (p *ncduPrinter) ProcessNode(path string, node *restic.Node) {
 	type NcduNode struct {
 		Name   string `json:"name"`
 		Asize  uint64 `json:"asize"`
@@ -168,30 +172,21 @@ func lsNodeNcdu(stdout io.Writer, depth *int, currentPath *string, path string,
 		Warnf("JSON encode failed: %v\n", err)
 	}
 
-	thisPath := filepath.Dir(path)
-	for thisPath != *currentPath {
-		*depth--
-		if *depth < 0 {
-			panic("cannot find suitable parent directory")
-		}
-		fmt.Fprintf(stdout, "\n%s]", strings.Repeat("  ", *depth))
-		*currentPath = filepath.Dir(*currentPath)
-	}
-
 	if node.Type == "dir" {
-		*currentPath = path
-		*depth++
-		fmt.Fprintf(stdout, ", [\n%s%s", strings.Repeat("  ", *depth), string(outJson))
+		p.depth++
+		fmt.Fprintf(p.out, ", [\n%s%s", strings.Repeat("  ", p.depth), string(outJson))
 	} else {
-		fmt.Fprintf(stdout, ",\n%s%s", strings.Repeat("  ", *depth), string(outJson))
+		fmt.Fprintf(p.out, ",\n%s%s", strings.Repeat("  ", p.depth), string(outJson))
 	}
 }
 
-func lsCloseNcdu(stdout io.Writer, depth *int) {
-	for *depth > 0 {
-		fmt.Fprintf(stdout, "%s]\n", strings.Repeat("  ", *depth))
-		*depth--
-	}
+func (p *ncduPrinter) LeaveDir(path string) {
+	p.depth--
+	fmt.Fprintf(p.out, "\n%s]", strings.Repeat("  ", p.depth))
+}
+
+func (p *ncduPrinter) Close() {
+	fmt.Fprint(p.out, "\n]\n")
 }
 
 func runLs(ctx context.Context, opts LsOptions, gopts GlobalOptions, args []string) error {
@@ -262,6 +257,8 @@ func runLs(ctx context.Context, opts LsOptions, gopts GlobalOptions, args []stri
 	var (
 		printSnapshot  func(sn *restic.Snapshot)
 		printNode      func(path string, node *restic.Node)
+		printLeaveNode func(path string)
+		printClose     func()
 	)
 
 	if gopts.JSON {
@@ -286,13 +283,13 @@ func runLs(ctx context.Context, opts LsOptions, gopts GlobalOptions, args []stri
 			}
 		}
 	} else if opts.Ncdu {
-		var depth int
-		var currentPath = "/"
-		printSnapshot = func(sn *restic.Snapshot) { lsSnapshotNcdu(globalOptions.stdout, &depth, sn) }
-		printNode = func(path string, node *restic.Node) {
-			lsNodeNcdu(globalOptions.stdout, &depth, &currentPath, path, node)
+		ncdu := &ncduPrinter{
+			out: globalOptions.stdout,
 		}
-		defer lsCloseNcdu(globalOptions.stdout, &depth)
+		printSnapshot = ncdu.ProcessSnapshot
+		printNode = ncdu.ProcessNode
+		printLeaveNode = ncdu.LeaveDir
+		printClose = ncdu.Close
 	} else {
 		printSnapshot = func(sn *restic.Snapshot) {
 			Verbosef("%v filtered by %v:\n", sn, dirs)
@@ -353,11 +350,20 @@ func runLs(ctx context.Context, opts LsOptions, gopts GlobalOptions, args []stri
 
 	err = walker.Walk(ctx, repo, *sn.Tree, walker.WalkVisitor{
 		ProcessNode: processNode,
+		LeaveDir: func(path string) {
+			if printLeaveNode != nil && withinDir(path) && path != "/" {
+				printLeaveNode(path)
+			}
+		},
 	})
 
 	if err != nil {
 		return err
 	}
 
+	if printClose != nil {
+		printClose()
+	}
+
 	return nil
 }