unix tree cmd

tree.go

package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
)

const (
	emptySpace   = "\t"
	middleItem   = "├───"
	continueLine = "│\t"
	lastItem     = "└───"
)

func getFormattedFileInfo(file os.FileInfo) string {
	var size string
	if file.Size() == 0 {
		size = "empty"
	} else {
		size = fmt.Sprint(file.Size()) + "b"
	}
	return fmt.Sprint(file.Name() + " (" + size + ")")
}

func tree(out io.Writer, path, indent string, withFiles bool) error {
	// 软链当成文件处理
	fileInfo, err := os.Lstat(path)
	if err != nil {
		return err
	}
	if fileInfo.Mode()&os.ModeSymlink != 0 {
		//fmt.Println(path, "is a symbolic link")
		if withFiles {
			fmt.Fprintln(out, getFormattedFileInfo(fileInfo))
		}
		return nil
	}

	file, err := os.Stat(path)
	if err != nil {
		return err
	}

	if !file.IsDir() {
		if withFiles {
			fmt.Fprintln(out, getFormattedFileInfo(file))
		}
		return nil
	}
	fmt.Fprintln(out, file.Name())

	subFiles, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	if !withFiles {
		dirs := make([]os.FileInfo, 0)
		for _, subFile := range subFiles {
			if subFile.IsDir() {
				dirs = append(dirs, subFile)
			}
		}
		subFiles = dirs
	}

	for i, subFile := range subFiles {
		newIndent := continueLine
		if i == len(subFiles)-1 {
			fmt.Fprintf(out, indent+lastItem)
			newIndent = emptySpace
		} else {
			fmt.Fprintf(out, indent+middleItem)
		}

		err := tree(out, path+"/"+subFile.Name(), indent+newIndent, withFiles)
		if err != nil {
			return err
		}
	}

	return nil
}

func main() {
	f, err := os.Create("./tmp")
	defer f.Close()
	if err != nil {
		panic(err)
	}

	out := f
	path, err := filepath.Abs(os.Args[1])
	err = tree(out, path, "", true)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}