Initial implementation of browse command. (#8)

Initial implementation of `browse` command.

.travis.yml

@@ -8,4 +8,8 @@ go:
 script: go test -v ./...
 
 env:
-  - GO111MODULE=on
+  - GO111MODULE=on
+
+branches:
+  only:
+    - master

cmd/browse/browse.go

@@ -0,0 +1,73 @@
+package browse
+
+import (
+	"errors"
+	"fmt"
+	"github.com/spf13/cobra"
+	"log"
+	"os"
+	"regexp"
+)
+
+type options struct {
+	tree bool
+	root string
+}
+
+func (o *options) initialise(cmd *cobra.Command, args []string) {
+	if len(args) != 0 {
+		r, _ := regexp.Compile("/$")
+		o.root = r.ReplaceAllString(args[0], "")
+	} else {
+		o.root = "."
+	}
+	if tree, _ := cmd.Flags().GetBool("tree"); tree {
+		o.tree = tree
+	}
+}
+
+func (o *options) validate() {
+	if err := o.validatePath(os.Stat(o.root)); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func (o *options) validatePath(info os.FileInfo, err error) error {
+	if os.IsNotExist(err) {
+		return errors.New(fmt.Sprintf("Directory \"%s\" does not exist", o.root))
+	}
+	return err
+}
+
+func (o *options) run() {
+	if !o.tree {
+		log.Fatal("Unknown display mode")
+		return
+	}
+	renderTree(buildFileTree(o.root))
+}
+
+var cmd = &cobra.Command{
+	Use:   "browse",
+	Short: "Interactively browse directories and files",
+}
+
+func NewInteractiveCmd() *cobra.Command {
+	o := options{}
+
+	cmd.Run = func(cmd *cobra.Command, args []string) {
+		o.initialise(cmd, args)
+		o.validate()
+		o.run()
+	}
+
+	cmd.Flags().BoolVarP(
+		&o.tree,
+		"tree",
+		"t",
+		true,
+		"browse the file tree",
+	)
+
+	return cmd
+}

cmd/browse/browse_test.go

@@ -0,0 +1,36 @@
+package browse
+
+import (
+	"errors"
+	"github.com/spf13/cobra"
+	"os"
+	"testing"
+)
+
+func TestInvalidPath(t *testing.T) {
+	cmd := cobra.Command{}
+	var args []string
+	o := options{}
+	o.initialise(&cmd, args)
+	var info os.FileInfo
+
+	err := o.validatePath(info, errors.New("something went wrong"))
+
+	if err == nil {
+		t.Errorf("Expected validation to fail when passing invalid path.")
+	}
+}
+
+func TestValidPath(t *testing.T) {
+	cmd := cobra.Command{}
+	var args []string
+	o := options{}
+	o.initialise(&cmd, args)
+	var info os.FileInfo
+
+	err := o.validatePath(info, nil)
+
+	if err != nil {
+		t.Errorf("Expected validation to pass when passing valid path.")
+	}
+}

cmd/browse/node.go

@@ -0,0 +1,41 @@
+package browse
+
+type node struct {
+	name     string
+	isDir    bool
+	size     int64
+	children []*node
+	parent   *node
+}
+
+func (n *node) addChild(c *node) {
+	n.children = append(n.children, c)
+	n.size += c.size
+}
+
+func (n *node) hasChild(name string) bool {
+	for _, c := range n.children {
+		if c.name == name {
+			return true
+		}
+	}
+	return false
+}
+
+func (n *node) getChild(name string) (*node, bool) {
+	for i, c := range n.children {
+		if c.name == name {
+			// Get a reference to array element directly, since range returns copies.
+			return n.children[i], true
+		}
+	}
+	return nil, false
+}
+
+func (n *node) recalculateSize() {
+	var s int64
+	for _, c := range n.children {
+		s += c.size
+	}
+	n.size = s
+}

cmd/browse/node_test.go

@@ -0,0 +1,72 @@
+package browse
+
+import (
+	"testing"
+)
+
+func TestCanAddAndRetrieveChild(t *testing.T) {
+	n := node{}
+	c1 := node{name: "C1", size: 100, isDir: false}
+	c2 := node{name: "C2", size: 200, isDir: false}
+	c3 := node{name: "C3", isDir: true}
+
+	n.addChild(&c1)
+	n.addChild(&c2)
+	n.addChild(&c3)
+
+	if len(n.children) != 3 {
+		t.Fatalf("Expected node to have 2 children, found %d", len(n.children))
+	}
+	if n.size != 300 {
+		t.Fatalf("Expected node to have size of %d, found %d", 300, n.size)
+	}
+	if !n.hasChild("C1") || !n.hasChild("C2") {
+		t.Fatalf("Expected node to have children C1 and C2")
+	}
+	if n.hasChild("C4") {
+		t.Fatalf("Expected node to not have child C4")
+	}
+	if c, ok := n.getChild("C1"); &c1 != c || !ok {
+		t.Fatalf("Expected node to have child C1")
+	}
+	if c, ok := n.getChild("C4"); c != nil || ok {
+		t.Fatalf("Expected node to not have child C4")
+	}
+}
+
+func TestCalculateSizeCorrectly(t *testing.T) {
+	n := node{name: "R"}
+	c1 := node{name: "C1", size: 100, isDir: false}
+	c2 := node{name: "C2", size: 200, isDir: false}
+	c3 := node{name: "C3", isDir: true}
+	c31 := node{name: "C31", size: 300, isDir: false}
+	c32 := node{name: "C32", size: 400, isDir: false}
+	c33 := node{name: "C33", isDir: true}
+	c331 := node{name: "C331", size: 500, isDir: false}
+	c332 := node{name: "C332", size: 600, isDir: false}
+	c333 := node{name: "C333", isDir: true}
+
+	n.addChild(&c1)
+	n.addChild(&c2)
+	n.addChild(&c3)
+	c3.addChild(&c31)
+	c3.addChild(&c32)
+	c3.addChild(&c33)
+	c33.addChild(&c331)
+	c33.addChild(&c332)
+	c33.addChild(&c333)
+
+	c33.recalculateSize()
+	c3.recalculateSize()
+	n.recalculateSize()
+
+	if c33.size != 1100 {
+		t.Fatalf("Expected child node c33 to have size of %d, found %d", 1100, c33.size)
+	}
+	if c3.size != 1800 {
+		t.Fatalf("Expected child node c3 to have size of %d, found %d", 1800, c33.size)
+	}
+	if n.size != 2100 {
+		t.Fatalf("Expected root node to have size of %d, found %d", 2100, c33.size)
+	}
+}

cmd/browse/tree.go

@@ -0,0 +1,143 @@
+package browse
+
+import (
+	"fmt"
+	"github.com/gdamore/tcell"
+	"github.com/rivo/tview"
+	"github.com/robinmitra/forest/formatter"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+)
+
+func buildNodesFromPath(n *node, path string, info os.FileInfo) {
+	nodeNames := strings.Split(path, "/")
+	currNodeName := nodeNames[0]
+	nestedNodeNames := nodeNames[1:]
+	// Last or trailing node
+	if len(nestedNodeNames) == 0 {
+		newNode := node{name: currNodeName, parent: n}
+		if info.IsDir() {
+			newNode.isDir = true
+		} else {
+			newNode.isDir = false
+			newNode.size = info.Size()
+		}
+		n.addChild(&newNode)
+	} else {
+		if existingNode, ok := n.getChild(currNodeName); ok {
+			buildNodesFromPath(existingNode, strings.Join(nestedNodeNames, "/"), info)
+			existingNode.recalculateSize()
+		} else {
+			newNode := node{name: currNodeName, isDir: true, parent: n}
+			buildNodesFromPath(&newNode, strings.Join(nestedNodeNames, "/"), info)
+			n.addChild(&newNode)
+		}
+	}
+}
+
+func processFile(node *node, rootPath string) filepath.WalkFunc {
+	return func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		name := info.Name()
+		if path == "." || path == rootPath {
+			node.name = name
+			return nil
+		}
+		if rootPath != "." {
+			buildNodesFromPath(node, strings.Replace(path, rootPath+"/", "", 1), info)
+		} else {
+			buildNodesFromPath(node, path, info)
+		}
+		return nil
+	}
+}
+
+func buildFileTree(root string) *node {
+	rootName := root
+	if root != "." {
+		path := strings.Split(root, "/")
+		rootName = path[len(path)-1]
+	}
+	rootNode := node{name: rootName, isDir: true}
+	if err := filepath.Walk(root, processFile(&rootNode, root)); err != nil {
+		log.Fatal(err)
+	}
+	return &rootNode
+}
+
+func renderTree(n *node) {
+	getNodeText := func(n *node) string {
+		return fmt.Sprintf("%s (%s, %d)", n.name, formatter.HumaniseStorage(n.size), len(n.children))
+	}
+
+	root := tview.NewTreeNode(getNodeText(n)).SetReference(n).SetColor(tcell.ColorRed)
+	tree := tview.NewTreeView().SetRoot(root).SetCurrentNode(root)
+
+	addChildren := func(n *tview.TreeNode) {
+		refNode := n.GetReference().(*node)
+		if len(refNode.children) > 0 {
+			for i, c := range refNode.children {
+				cNode := refNode.children[i]
+				childNode := tview.NewTreeNode(getNodeText(cNode)).SetReference(cNode)
+				if c.isDir {
+					childNode.SetColor(tcell.ColorGreen)
+				}
+				n.AddChild(childNode)
+			}
+		}
+	}
+
+	addChildren(root)
+
+	tree.SetSelectedFunc(func(n *tview.TreeNode) {
+		refNode := n.GetReference()
+		if refNode == nil {
+			// Selecting the root node does nothing.
+			return
+		}
+		children := n.GetChildren()
+		if len(children) == 0 {
+			// Load and show files in this directory.
+			addChildren(n)
+		} else {
+			// Collapse if visible, expand if collapsed.
+			n.SetExpanded(!n.IsExpanded())
+		}
+	})
+
+	app := tview.NewApplication().SetRoot(tree, true)
+
+	tree.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
+		if event.Key() == tcell.KeyRune && event.Rune() == 'o' {
+			return tcell.NewEventKey(tcell.KeyEnter, 0, tcell.ModNone)
+		}
+		return event
+	})
+
+	if err := app.Run(); err != nil {
+		panic(err)
+	}
+}
+
+func debugTree(n *node, spacer string) {
+	var b strings.Builder
+	fmt.Fprintf(&b, "%s %s", spacer, n.name)
+	fmt.Fprintf(&b, " (")
+	fmt.Fprintf(&b, "dir: %t", n.isDir)
+	if n.parent != nil {
+		fmt.Fprintf(&b, ", parent: %s", n.parent.name)
+	}
+	fmt.Fprintf(&b, ", size: %d", n.size)
+	fmt.Fprintf(&b, ")")
+	fmt.Println(b.String())
+
+	if n.isDir && len(n.children) > 0 {
+		for i, _ := range n.children {
+			debugTree(n.children[i], fmt.Sprintf("%s-", spacer))
+		}
+	}
+}

cmd/root.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"github.com/robinmitra/forest/cmd/analyse"
+	"github.com/robinmitra/forest/cmd/browse"
 	"github.com/robinmitra/forest/cmd/version"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -47,6 +48,7 @@ func NewRootCmd() *cobra.Command {
 
 	cmd.AddCommand(analyse.NewAnalyseCmd())
 	cmd.AddCommand(version.NewVersionCmd(VERSION))
+	cmd.AddCommand(browse.NewInteractiveCmd())
 
 	return cmd
 }