Get folder size iteratively in go

go-tree.go

// https://github.com/alexavlonitis
// Get folder size iteratively in go.
// The same can be achieved using the file.WalkDir() function that recursively walks a directory.

package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
)

func main() {
	basePath := "."

	tree := buildTree(basePath)
	parseTree(tree)
}

// Build a hierarchy tree based on the given path and its subfolders
// Does not use recursion, custom walk.
// Returns a tree struct
func buildTree(path string) Tree {
	var totalSize int64
	node := Node{value: path, depth: 0}
	tree := Tree{root: &node}

	nodePathMap := make(map[string]*Node)
	nodePathMap[path] = &node

	queue := []string{path}

	for len(queue) > 0 {
		folderPath := queue[0]
		parentNode := nodePathMap[folderPath]

		files, err := ioutil.ReadDir(folderPath)
		if err != nil {
			log.Fatal(err)
		}

		for _, file := range files {
			fullPath := folderPath + "/" + file.Name()

			if file.IsDir() && !(file.Mode()&os.ModeSymlink == os.ModeSymlink) {
				childNode := Node{value: fullPath, parent: parentNode, depth: parentNode.depth + 1}
				parentNode.children = append(parentNode.children, &childNode)
				nodePathMap[fullPath] = &childNode

				queue = append(queue, fullPath)
			} else {
				s := file.Size()
				totalSize += s
				parentNode.size += s
			}
		}
		queue = queue[1:]
	}

	node.size = totalSize

	return tree
}

// BreadthFirst traversal
// Prints the hierarchy and size of each node and the total size of the base path
func parseTree(tree Tree) {
	queue := []Node{*tree.root}

	for len(queue) > 0 {
		node := queue[0]
		for _, n := range node.children {
			queue = append(queue, *n)
		}
		queue = queue[1:]

		fmt.Println(node.value, "["+ByteCountSI(node.size)+"]")
	}

	fmt.Println("Total value:", tree.root.value, ByteCountSI(tree.root.size))
}

// Convert bytes to KB, MB etc...
// Returns the size in a string representation
func ByteCountSI(b int64) string {
	const unit = 1000
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "kMGTPE"[exp])
}

type Tree struct {
	root *Node
}
type Node struct {
	parent   *Node
	children []*Node
	value    string
	size     int64
	depth    int
}