list directory contents like tree command in Go

tree.go

package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"

	colorable "github.com/mattn/go-colorable"
)

func main() {
	os.Exit(run())
}

func run() int {
	flag.Parse()

	flag.Usage = func() {
		fmt.Println("Usage: tree directory")
	}
	if flag.NArg() != 1 {
		flag.Usage()
		return 1
	}

	dir := flag.Arg(0)
	p := newTreePrinter()
	err := p.showTree(dir)
	if err != nil {
		fmt.Println(err)
		return 1
	}
	return 0
}

type treePrinter struct {
	out       io.Writer
	fileCount int
	dirCount  int
}

func newTreePrinter() *treePrinter {
	return &treePrinter{
		out: colorable.NewColorableStdout(),
	}
}

func (p *treePrinter) Println(a ...interface{}) (n int, err error) {
	return fmt.Fprintln(p.out, a...)
}

func (p *treePrinter) Printf(format string, a ...interface{}) (n int, err error) {
	return fmt.Fprintf(p.out, format, a...)
}

func (p *treePrinter) showTree(dir string) error {
	p.Println(colorStr(blue, dir))
	err := p.showTreeHelper("", dir)
	if err != nil {
		return err
	}

	p.Println()
	p.Printf("%d dirctories, %d files\n", p.dirCount, p.fileCount)
	return nil
}

func (p *treePrinter) showTreeHelper(prefix, dir string) error {
	infos, err := readDir(dir)
	if err != nil {
		return err
	}
	sortFileInfosByName(infos)

	for i, info := range infos {
		if strings.HasPrefix(info.Name(), ".") {
			continue
		}
		var branch string
		if i < len(infos)-1 {
			branch = "├── "
		} else {
			branch = "└── "
		}
		var name string
		if info.IsDir() {
			name = colorStr(blue, info.Name())
		} else if isExecutable(info.Mode()) {
			name = colorStr(green, info.Name())
		} else {
			name = info.Name()
		}
		p.Printf("%s%s%s\n", prefix, branch, name)
		if info.IsDir() {
			var nextPrefix string
			if i < len(infos)-1 {
				nextPrefix = prefix + "│   "
			} else {
				nextPrefix = prefix + "    "
			}
			err = p.showTreeHelper(nextPrefix, filepath.Join(dir, info.Name()))
			if err != nil {
				return err
			}
			p.dirCount++
		} else {
			p.fileCount++
		}
	}

	return nil
}

func isExecutable(mode os.FileMode) bool {
	return mode&0100 != 0
}

func readDir(path string) ([]os.FileInfo, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	infos, err := file.Readdir(0)
	if err != nil {
		return nil, err
	}
	return infos, nil
}

type fileInfosByName []os.FileInfo

func (a fileInfosByName) Len() int      { return len(a) }
func (a fileInfosByName) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a fileInfosByName) Less(i, j int) bool {
	return strings.ToLower(a[i].Name()) < strings.ToLower(a[j].Name())
}

func sortFileInfosByName(infos []os.FileInfo) {
	sort.Sort(fileInfosByName(infos))
}

const (
	blue  = "34"
	green = "32"
)

func colorStr(color, s string) string {
	return fmt.Sprintf("\x1b[%sm%s\x1b[0m", color, s)
}