Add output flag and simple tree option

print/flat.go

@@ -0,0 +1,33 @@
+package print
+
+import (
+	"fmt"
+	"git-get/git"
+	"path/filepath"
+	"strings"
+)
+
+type FlatPrinter struct{}
+
+func (p *FlatPrinter) Print(root string, repos []*git.Repo) string {
+	val := root
+
+	for _, repo := range repos {
+		path := strings.TrimPrefix(repo.Path, root)
+		path = strings.Trim(path, string(filepath.Separator))
+
+		val += fmt.Sprintf("\n%s %s", path, printWorktreeStatus(repo))
+
+		for _, branch := range repo.Status.Branches {
+			// Don't print the status of the current branch. It was already printed above.
+			if branch.Name == repo.Status.CurrentBranch {
+				continue
+			}
+
+			indent := strings.Repeat(" ", len(path))
+			val += fmt.Sprintf("\n%s %s", indent, printBranchStatus(branch))
+		}
+	}
+
+	return val
+}

print/print.go

@@ -3,7 +3,6 @@ package print
 import (
 	"fmt"
 	"git-get/git"
-	"path/filepath"
 	"strings"
 )
 
@@ -11,35 +10,6 @@ type Printer interface {
 	Print(root string, repos []*git.Repo) string
 }
 
-type FlatPrinter struct{}
-
-func NewFlatPrinter() *FlatPrinter {
-	return &FlatPrinter{}
-}
-
-func (p *FlatPrinter) Print(root string, repos []*git.Repo) string {
-	val := root
-
-	for _, repo := range repos {
-		path := strings.TrimPrefix(repo.Path, root)
-		path = strings.Trim(path, string(filepath.Separator))
-
-		val += fmt.Sprintf("\n%s %s", path, renderWorktreeStatus(repo))
-
-		for _, branch := range repo.Status.Branches {
-			// Don't print the status of the current branch. It was already printed above.
-			if branch.Name == repo.Status.CurrentBranch {
-				continue
-			}
-
-			indent := strings.Repeat(" ", len(path))
-			val += fmt.Sprintf("\n%s %s", indent, renderBranchStatus(branch))
-		}
-	}
-
-	return val
-}
-
 const (
 	ColorRed    = "\033[1;31m%s\033[0m"
 	ColorGreen  = "\033[1;32m%s\033[0m"
@@ -47,7 +17,7 @@ const (
 	ColorYellow = "\033[1;33m%s\033[0m"
 )
 
-func renderWorktreeStatus(repo *git.Repo) string {
+func printWorktreeStatus(repo *git.Repo) string {
 	clean := true
 	var status []string
 
@@ -56,7 +26,7 @@ func renderWorktreeStatus(repo *git.Repo) string {
 	if current := repo.CurrentBranchStatus(); current == nil {
 		status = append(status, fmt.Sprintf(ColorYellow, repo.Status.CurrentBranch))
 	} else {
-		status = append(status, renderBranchStatus(current))
+		status = append(status, printBranchStatus(current))
 	}
 
 	// TODO: this is ugly
@@ -85,7 +55,7 @@ func renderWorktreeStatus(repo *git.Repo) string {
 	return strings.Join(status, " ")
 }
 
-func renderBranchStatus(branch *git.BranchStatus) string {
+func printBranchStatus(branch *git.BranchStatus) string {
 	// ok indicates that the branch has upstream and is not ahead or behind it
 	ok := true
 	var status []string

print/tree.go

@@ -4,12 +4,39 @@ import (
 	"git-get/git"
 	"path/filepath"
 	"strings"
+
+	"github.com/xlab/treeprint"
 )
 
-// Node represents a node in a repos tree
+type SimpleTreePrinter struct{}
+
+type SmartTreePrinter struct {
+	// length is the size (number of chars) of the currently processed line.
+	// It's used to correctly indent the lines with branches status.
+	length int
+}
+
+func (p *SmartTreePrinter) Print(root string, repos []*git.Repo) string {
+	tree := BuildTree(root, repos)
+
+	return p.PrintSmartTree(tree)
+}
+
+func (p *SimpleTreePrinter) Print(root string, repos []*git.Repo) string {
+	tree := BuildTree(root, repos)
+
+	tp := treeprint.New()
+	tp.SetValue(root)
+
+	p.PrintSimpleTree(tree, tp)
+
+	return tp.String()
+}
+
+// Node represents a node (ie. path fragment) in a repos tree.
 type Node struct {
 	val      string
-	depth    int // depth is a nesting depth used when rendering a tree, not an depth level of a node inside the tree
+	depth    int // depth is a nesting depth used when rendering a smart tree, not a depth level of a tree node.
 	parent   *Node
 	children []*Node
 	repo     *git.Repo
@@ -87,7 +114,7 @@ func BuildTree(root string, repos []*git.Repo) *Node {
 	return tree
 }
 
-// RenderSmartTree returns a string representation of repos tree.
+// PrintSmartTree returns a string representation of repos tree.
 // It's "smart" because it automatically folds branches which only have a single child and indents branches with many children.
 //
 // It recursively traverses the tree and prints its nodes.
@@ -109,7 +136,7 @@ func BuildTree(root string, repos []*git.Repo) *Node {
 //           repo2
 //       another/repo
 //
-func RenderSmartTree(node *Node) string {
+func (p *SmartTreePrinter) PrintSmartTree(node *Node) string {
 	if node.children == nil {
 		// If node is a leaf, print repo name and its status and finish processing this node.
 		value := node.val
@@ -120,17 +147,17 @@ func RenderSmartTree(node *Node) string {
 			return value
 		}
 
-		value += " " + renderWorktreeStatus(node.repo)
+		value += " " + printWorktreeStatus(node.repo)
 
 		// Print the status of each branch on a new line, indented to match the position of the current branch name.
-		indent := "\n" + strings.Repeat(" ", length+len(node.val))
+		indent := "\n" + strings.Repeat(" ", p.length+len(node.val))
 		for _, branch := range node.repo.Status.Branches {
 			// Don't print the status of the current branch. It was already printed above.
 			if branch.Name == node.repo.Status.CurrentBranch {
 				continue
 			}
 
-			value += indent + renderBranchStatus(branch)
+			value += indent + printBranchStatus(branch)
 		}
 
 		return value
@@ -148,24 +175,40 @@ func RenderSmartTree(node *Node) string {
 		// node's path has multiple levels above.
 		node.depth = node.parent.depth
 
-		length += len(val)
+		p.length += len(val)
 	} else {
 		// If node has multiple children, print each of them on a new line
 		// and indent them once relative to the parent
 		node.depth = node.parent.depth + 1
 		shift = "\n" + strings.Repeat("    ", node.depth)
-		length = 0
+		p.length = 0
 	}
 
 	for _, child := range node.children {
-		length += len(shift)
-		val += shift + RenderSmartTree(child)
-		length = 0
+		p.length += len(shift)
+		val += shift + p.PrintSmartTree(child)
+		p.length = 0
 	}
 
 	return val
 }
 
-// lenght is the size (number of chars) of the currently processed line.
-// It's used to correctly indent the lines with branches status.
-var length int
+func (p *SimpleTreePrinter) PrintSimpleTree(node *Node, tp treeprint.Tree) {
+	if node.children == nil {
+		tp.SetValue(node.val + " " + printWorktreeStatus(node.repo))
+
+		for _, branch := range node.repo.Status.Branches {
+			// Don't print the status of the current branch. It was already printed above.
+			if branch.Name == node.repo.Status.CurrentBranch {
+				continue
+			}
+			tp.AddNode(printBranchStatus(branch))
+		}
+
+	}
+
+	for _, child := range node.children {
+		branch := tp.AddBranch(child.val)
+		p.PrintSimpleTree(child, branch)
+	}
+}

print/tree_test.go

@@ -96,9 +96,9 @@ gitlab.com/
 			repos = append(repos, git.NewRepo(nil, path)) //&Repo{path: path})
 		}
 
-		tree := BuildTree("root", repos)
+		printer := SmartTreePrinter{}
 		// Leading and trailing newlines are added to test cases for readability. We also need to add them to the rendering result.
-		got := fmt.Sprintf("\n%s\n", RenderSmartTree(tree))
+		got := fmt.Sprintf("\n%s\n", printer.Print("root", repos))
 
 		// Rendered tree uses spaces for indentation but the test cases use tabs.
 		if got != strings.ReplaceAll(test.want, "\t", "    ") {