github.com/while1malloc0/docdir@v0.0.0-20220830001304-722ec0f2cf3a/dirtree/dirtree.go (about)

     1  // package dirtree implements a tree structure of directories
     2  package dirtree
     3  
     4  import (
     5  	"bufio"
     6  	"fmt"
     7  	"io"
     8  	"os"
     9  	"path/filepath"
    10  	"sort"
    11  	"strings"
    12  )
    13  
    14  const DefaultDescriptionFile = "DESCRIPTION"
    15  
    16  type Node struct {
    17  	Name        string
    18  	Description string
    19  	Children    []*Node
    20  }
    21  
    22  func New(path string, skipMissing bool) (*Node, error) {
    23  	fi, err := os.Stat(path)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  	if !fi.IsDir() {
    28  		return nil, nil
    29  	}
    30  
    31  	n := Node{Name: fi.Name(), Children: []*Node{}}
    32  	descriptionData, err := os.ReadFile(filepath.Join(path, DefaultDescriptionFile))
    33  	if os.IsNotExist(err) && skipMissing {
    34  		return nil, nil
    35  	}
    36  	if err == nil {
    37  		n.Description = string(descriptionData)
    38  	}
    39  
    40  	dir, err := os.Open(path)
    41  	if err != nil {
    42  		return nil, err
    43  	}
    44  	names, err := dir.Readdirnames(-1)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	_ = dir.Close()
    49  
    50  	names = filterNonDirs(path, names)
    51  	sort.Strings(names)
    52  	for _, name := range names {
    53  		child, err := New(filepath.Join(path, name), skipMissing)
    54  		if err != nil {
    55  			return nil, err
    56  		}
    57  		// TODO [jturner 2022-08-29]: don't love this
    58  		if child != nil {
    59  			n.Children = append(n.Children, child)
    60  		}
    61  	}
    62  
    63  	return &n, nil
    64  }
    65  
    66  func (n *Node) String() string {
    67  	var s strings.Builder
    68  	visit(&s, n, "")
    69  	return align(s.String())
    70  }
    71  
    72  // based on https://github.com/campoy/tools/tree
    73  func visit(w io.Writer, node *Node, prefix string) error {
    74  	s := node.Name
    75  	if node.Description != "" {
    76  		s = fmt.Sprintf("%s # %s", node.Name, node.Description)
    77  	}
    78  	fmt.Fprintln(w, s)
    79  	padding := "│   "
    80  	for i, child := range node.Children {
    81  		if i == len(node.Children)-1 {
    82  			fmt.Fprintf(w, prefix+"└── ")
    83  			padding = "    "
    84  		} else {
    85  			fmt.Fprintf(w, prefix+"├── ")
    86  		}
    87  		err := visit(w, child, prefix+padding)
    88  		if err != nil {
    89  			return err
    90  		}
    91  	}
    92  	return nil
    93  }
    94  
    95  func align(in string) string {
    96  	var out strings.Builder
    97  
    98  	longest := findLongest(in)
    99  	// longest name+indent pair + 2 spaces
   100  	lengthToMeet := longest + 2
   101  	scanner := bufio.NewScanner(strings.NewReader(in))
   102  	for scanner.Scan() {
   103  		line := scanner.Text()
   104  		parts := strings.Split(line, "#")
   105  		if len(parts) == 1 {
   106  			// no description provided
   107  			fmt.Fprintln(&out, line)
   108  			continue
   109  		}
   110  		nameAndIndent := strings.TrimRight(parts[0], " ")
   111  		description := strings.TrimSpace(parts[1])
   112  		paddingNeeded := lengthToMeet - normalizedLen(nameAndIndent)
   113  		padding := strings.Repeat(" ", paddingNeeded)
   114  		fmt.Fprintf(&out, "%s%s# %s\n", nameAndIndent, padding, description)
   115  	}
   116  	return out.String()
   117  }
   118  
   119  func filterNonDirs(path string, candidates []string) []string {
   120  	var dirs []string
   121  	for _, candidate := range candidates {
   122  		fi, _ := os.Stat(filepath.Join(path, candidate))
   123  		if fi.IsDir() {
   124  			dirs = append(dirs, candidate)
   125  		}
   126  	}
   127  	return dirs
   128  }
   129  
   130  // return the length of a string from a visual perspective, e.g. each character
   131  // is one space, regardless of unicode
   132  func normalizedLen(s string) int {
   133  	s = strings.Replace(s, "├──", "   ", 1)
   134  	s = strings.Replace(s, "└──", "   ", 1)
   135  	return len(s)
   136  }
   137  
   138  // returns the length of the longest string after stripping descriptions and
   139  // normalizing string len
   140  func findLongest(s string) int {
   141  	scanner := bufio.NewScanner(strings.NewReader(s))
   142  	longest := -1
   143  	for scanner.Scan() {
   144  		line := scanner.Text()
   145  		parts := strings.Split(line, "#")
   146  		nameAndIndent := strings.TrimRight(parts[0], " ")
   147  		visualLen := normalizedLen(nameAndIndent)
   148  		if visualLen > longest {
   149  			longest = visualLen
   150  		}
   151  	}
   152  	return longest
   153  }