treeprint.go

Documentation: github.com/xlab/treeprint

     1  // Package treeprint provides a simple ASCII tree composing tool.
     2  package treeprint
     3  
     4  import (
     5  	"bytes"
     6  	"fmt"
     7  	"io"
     8  	"reflect"
     9  	"strings"
    10  )
    11  
    12  // Value defines any value
    13  type Value interface{}
    14  
    15  // MetaValue defines any meta value
    16  type MetaValue interface{}
    17  
    18  // NodeVisitor function type for iterating over nodes
    19  type NodeVisitor func(item *Node)
    20  
    21  // Tree represents a tree structure with leaf-nodes and branch-nodes.
    22  type Tree interface {
    23  	// AddNode adds a new Node to a branch.
    24  	AddNode(v Value) Tree
    25  	// AddMetaNode adds a new Node with meta value provided to a branch.
    26  	AddMetaNode(meta MetaValue, v Value) Tree
    27  	// AddBranch adds a new branch Node (a level deeper).
    28  	AddBranch(v Value) Tree
    29  	// AddMetaBranch adds a new branch Node (a level deeper) with meta value provided.
    30  	AddMetaBranch(meta MetaValue, v Value) Tree
    31  	// Branch converts a leaf-Node to a branch-Node,
    32  	// applying this on a branch-Node does no effect.
    33  	Branch() Tree
    34  	// FindByMeta finds a Node whose meta value matches the provided one by reflect.DeepEqual,
    35  	// returns nil if not found.
    36  	FindByMeta(meta MetaValue) Tree
    37  	// FindByValue finds a Node whose value matches the provided one by reflect.DeepEqual,
    38  	// returns nil if not found.
    39  	FindByValue(value Value) Tree
    40  	//  returns the last Node of a tree
    41  	FindLastNode() Tree
    42  	// String renders the tree or subtree as a string.
    43  	String() string
    44  	// Bytes renders the tree or subtree as byteslice.
    45  	Bytes() []byte
    46  
    47  	SetValue(value Value)
    48  	SetMetaValue(meta MetaValue)
    49  
    50  	// VisitAll iterates over the tree, branches and nodes.
    51  	// If need to iterate over the whole tree, use the root Node.
    52  	// Note this method uses a breadth-first approach.
    53  	VisitAll(fn NodeVisitor)
    54  }
    55  
    56  type Node struct {
    57  	Root  *Node
    58  	Meta  MetaValue
    59  	Value Value
    60  	Nodes []*Node
    61  }
    62  
    63  func (n *Node) FindLastNode() Tree {
    64  	ns := n.Nodes
    65  	if len(ns) == 0 {
    66  		return nil
    67  	}
    68  	return ns[len(ns)-1]
    69  }
    70  
    71  func (n *Node) AddNode(v Value) Tree {
    72  	n.Nodes = append(n.Nodes, &Node{
    73  		Root:  n,
    74  		Value: v,
    75  	})
    76  	return n
    77  }
    78  
    79  func (n *Node) AddMetaNode(meta MetaValue, v Value) Tree {
    80  	n.Nodes = append(n.Nodes, &Node{
    81  		Root:  n,
    82  		Meta:  meta,
    83  		Value: v,
    84  	})
    85  	return n
    86  }
    87  
    88  func (n *Node) AddBranch(v Value) Tree {
    89  	branch := &Node{
    90  		Root:  n,
    91  		Value: v,
    92  	}
    93  	n.Nodes = append(n.Nodes, branch)
    94  	return branch
    95  }
    96  
    97  func (n *Node) AddMetaBranch(meta MetaValue, v Value) Tree {
    98  	branch := &Node{
    99  		Root:  n,
   100  		Meta:  meta,
   101  		Value: v,
   102  	}
   103  	n.Nodes = append(n.Nodes, branch)
   104  	return branch
   105  }
   106  
   107  func (n *Node) Branch() Tree {
   108  	n.Root = nil
   109  	return n
   110  }
   111  
   112  func (n *Node) FindByMeta(meta MetaValue) Tree {
   113  	for _, node := range n.Nodes {
   114  		if reflect.DeepEqual(node.Meta, meta) {
   115  			return node
   116  		}
   117  		if v := node.FindByMeta(meta); v != nil {
   118  			return v
   119  		}
   120  	}
   121  	return nil
   122  }
   123  
   124  func (n *Node) FindByValue(value Value) Tree {
   125  	for _, node := range n.Nodes {
   126  		if reflect.DeepEqual(node.Value, value) {
   127  			return node
   128  		}
   129  		if v := node.FindByMeta(value); v != nil {
   130  			return v
   131  		}
   132  	}
   133  	return nil
   134  }
   135  
   136  func (n *Node) Bytes() []byte {
   137  	buf := new(bytes.Buffer)
   138  	level := 0
   139  	var levelsEnded []int
   140  	if n.Root == nil {
   141  		if n.Meta != nil {
   142  			buf.WriteString(fmt.Sprintf("[%v]  %v", n.Meta, n.Value))
   143  		} else {
   144  			buf.WriteString(fmt.Sprintf("%v", n.Value))
   145  		}
   146  		buf.WriteByte('\n')
   147  	} else {
   148  		edge := EdgeTypeMid
   149  		if len(n.Nodes) == 0 {
   150  			edge = EdgeTypeEnd
   151  			levelsEnded = append(levelsEnded, level)
   152  		}
   153  		printValues(buf, 0, levelsEnded, edge, n)
   154  	}
   155  	if len(n.Nodes) > 0 {
   156  		printNodes(buf, level, levelsEnded, n.Nodes)
   157  	}
   158  	return buf.Bytes()
   159  }
   160  
   161  func (n *Node) String() string {
   162  	return string(n.Bytes())
   163  }
   164  
   165  func (n *Node) SetValue(value Value) {
   166  	n.Value = value
   167  }
   168  
   169  func (n *Node) SetMetaValue(meta MetaValue) {
   170  	n.Meta = meta
   171  }
   172  
   173  func (n *Node) VisitAll(fn NodeVisitor) {
   174  	for _, node := range n.Nodes {
   175  		fn(node)
   176  
   177  		if len(node.Nodes) > 0 {
   178  			node.VisitAll(fn)
   179  			continue
   180  		}
   181  	}
   182  }
   183  
   184  func printNodes(wr io.Writer,
   185  	level int, levelsEnded []int, nodes []*Node) {
   186  
   187  	for i, node := range nodes {
   188  		edge := EdgeTypeMid
   189  		if i == len(nodes)-1 {
   190  			levelsEnded = append(levelsEnded, level)
   191  			edge = EdgeTypeEnd
   192  		}
   193  		printValues(wr, level, levelsEnded, edge, node)
   194  		if len(node.Nodes) > 0 {
   195  			printNodes(wr, level+1, levelsEnded, node.Nodes)
   196  		}
   197  	}
   198  }
   199  
   200  func printValues(wr io.Writer,
   201  	level int, levelsEnded []int, edge EdgeType, node *Node) {
   202  
   203  	for i := 0; i < level; i++ {
   204  		if isEnded(levelsEnded, i) {
   205  			fmt.Fprint(wr, strings.Repeat(" ", IndentSize+1))
   206  			continue
   207  		}
   208  		fmt.Fprintf(wr, "%s%s", EdgeTypeLink, strings.Repeat(" ", IndentSize))
   209  	}
   210  
   211  	val := renderValue(level, node)
   212  	meta := node.Meta
   213  
   214  	if meta != nil {
   215  		fmt.Fprintf(wr, "%s [%v]  %v\n", edge, meta, val)
   216  		return
   217  	}
   218  	fmt.Fprintf(wr, "%s %v\n", edge, val)
   219  }
   220  
   221  func isEnded(levelsEnded []int, level int) bool {
   222  	for _, l := range levelsEnded {
   223  		if l == level {
   224  			return true
   225  		}
   226  	}
   227  	return false
   228  }
   229  
   230  func renderValue(level int, node *Node) Value {
   231  	lines := strings.Split(fmt.Sprintf("%v", node.Value), "\n")
   232  
   233  	// If value does not contain multiple lines, return itself.
   234  	if len(lines) < 2 {
   235  		return node.Value
   236  	}
   237  
   238  	// If value contains multiple lines,
   239  	// generate a padding and prefix each line with it.
   240  	pad := padding(level, node)
   241  
   242  	for i := 1; i < len(lines); i++ {
   243  		lines[i] = fmt.Sprintf("%s%s", pad, lines[i])
   244  	}
   245  
   246  	return strings.Join(lines, "\n")
   247  }
   248  
   249  // padding returns a padding for the multiline values with correctly placed link edges.
   250  // It is generated by traversing the tree upwards (from leaf to the root of the tree)
   251  // and, on each level, checking if the Node the last one of its siblings.
   252  // If a Node is the last one, the padding on that level should be empty (there's nothing to link to below it).
   253  // If a Node is not the last one, the padding on that level should be the link edge so the sibling below is correctly connected.
   254  func padding(level int, node *Node) string {
   255  	links := make([]string, level+1)
   256  
   257  	for node.Root != nil {
   258  		if isLast(node) {
   259  			links[level] = strings.Repeat(" ", IndentSize+1)
   260  		} else {
   261  			links[level] = fmt.Sprintf("%s%s", EdgeTypeLink, strings.Repeat(" ", IndentSize))
   262  		}
   263  		level--
   264  		node = node.Root
   265  	}
   266  
   267  	return strings.Join(links, "")
   268  }
   269  
   270  // isLast checks if the Node is the last one in the slice of its parent children
   271  func isLast(n *Node) bool {
   272  	return n == n.Root.FindLastNode()
   273  }
   274  
   275  type EdgeType string
   276  
   277  var (
   278  	EdgeTypeLink EdgeType = "│"
   279  	EdgeTypeMid  EdgeType = "├──"
   280  	EdgeTypeEnd  EdgeType = "└──"
   281  )
   282  
   283  // IndentSize is the number of spaces per tree level.
   284  var IndentSize = 3
   285  
   286  // New Generates new tree
   287  func New() Tree {
   288  	return &Node{Value: "."}
   289  }
   290  
   291  // NewWithRoot Generates new tree with the given root value
   292  func NewWithRoot(root Value) Tree {
   293  	return &Node{Value: root}
   294  }
   295
View as plain text