k3s/vendor/github.com/xlab/treeprint/treeprint.go

216 lines
4.4 KiB
Go
Raw Blame History

This file contains invisible Unicode characters!

This file contains invisible Unicode characters that may be processed differently from what appears below. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to reveal hidden characters.

// Package treeprint provides a simple ASCII tree composing tool.
package treeprint
import (
"bytes"
"fmt"
"io"
"reflect"
)
type Value interface{}
type MetaValue interface{}
// Tree represents a tree structure with leaf-nodes and branch-nodes.
type Tree interface {
// AddNode adds a new node to a branch.
AddNode(v Value) Tree
// AddMetaNode adds a new node with meta value provided to a branch.
AddMetaNode(meta MetaValue, v Value) Tree
// AddBranch adds a new branch node (a level deeper).
AddBranch(v Value) Tree
// AddMetaBranch adds a new branch node (a level deeper) with meta value provided.
AddMetaBranch(meta MetaValue, v Value) Tree
// Branch converts a leaf-node to a branch-node,
// applying this on a branch-node does no effect.
Branch() Tree
// FindByMeta finds a node whose meta value matches the provided one by reflect.DeepEqual,
// returns nil if not found.
FindByMeta(meta MetaValue) Tree
// FindByValue finds a node whose value matches the provided one by reflect.DeepEqual,
// returns nil if not found.
FindByValue(value Value) Tree
// returns the last node of a tree
FindLastNode() Tree
// String renders the tree or subtree as a string.
String() string
// Bytes renders the tree or subtree as byteslice.
Bytes() []byte
SetValue(value Value)
SetMetaValue(meta MetaValue)
}
type node struct {
Root *node
Meta MetaValue
Value Value
Nodes []*node
}
func (n *node) FindLastNode() Tree {
ns := n.Nodes
n = ns[len(ns)-1]
return n
}
func (n *node) AddNode(v Value) Tree {
n.Nodes = append(n.Nodes, &node{
Root: n,
Value: v,
})
if n.Root != nil {
return n.Root
}
return n
}
func (n *node) AddMetaNode(meta MetaValue, v Value) Tree {
n.Nodes = append(n.Nodes, &node{
Root: n,
Meta: meta,
Value: v,
})
if n.Root != nil {
return n.Root
}
return n
}
func (n *node) AddBranch(v Value) Tree {
branch := &node{
Value: v,
}
n.Nodes = append(n.Nodes, branch)
return branch
}
func (n *node) AddMetaBranch(meta MetaValue, v Value) Tree {
branch := &node{
Meta: meta,
Value: v,
}
n.Nodes = append(n.Nodes, branch)
return branch
}
func (n *node) Branch() Tree {
n.Root = nil
return n
}
func (n *node) FindByMeta(meta MetaValue) Tree {
for _, node := range n.Nodes {
if reflect.DeepEqual(node.Meta, meta) {
return node
}
if v := node.FindByMeta(meta); v != nil {
return v
}
}
return nil
}
func (n *node) FindByValue(value Value) Tree {
for _, node := range n.Nodes {
if reflect.DeepEqual(node.Value, value) {
return node
}
if v := node.FindByMeta(value); v != nil {
return v
}
}
return nil
}
func (n *node) Bytes() []byte {
buf := new(bytes.Buffer)
level := 0
var levelsEnded []int
if n.Root == nil {
if n.Meta != nil {
buf.WriteString(fmt.Sprintf("[%v] %v", n.Meta, n.Value))
} else {
buf.WriteString(fmt.Sprintf("%v",n.Value))
}
buf.WriteByte('\n')
} else {
edge := EdgeTypeMid
if len(n.Nodes) == 0 {
edge = EdgeTypeEnd
levelsEnded = append(levelsEnded, level)
}
printValues(buf, 0, levelsEnded, edge, n.Meta, n.Value)
}
if len(n.Nodes) > 0 {
printNodes(buf, level, levelsEnded, n.Nodes)
}
return buf.Bytes()
}
func (n *node) String() string {
return string(n.Bytes())
}
func (n *node) SetValue(value Value){
n.Value = value
}
func (n *node) SetMetaValue(meta MetaValue){
n.Meta = meta
}
func printNodes(wr io.Writer,
level int, levelsEnded []int, nodes []*node) {
for i, node := range nodes {
edge := EdgeTypeMid
if i == len(nodes)-1 {
levelsEnded = append(levelsEnded, level)
edge = EdgeTypeEnd
}
printValues(wr, level, levelsEnded, edge, node.Meta, node.Value)
if len(node.Nodes) > 0 {
printNodes(wr, level+1, levelsEnded, node.Nodes)
}
}
}
func printValues(wr io.Writer,
level int, levelsEnded []int, edge EdgeType, meta MetaValue, val Value) {
for i := 0; i < level; i++ {
if isEnded(levelsEnded, i) {
fmt.Fprint(wr, " ")
continue
}
fmt.Fprintf(wr, "%s   ", EdgeTypeLink)
}
if meta != nil {
fmt.Fprintf(wr, "%s [%v] %v\n", edge, meta, val)
return
}
fmt.Fprintf(wr, "%s %v\n", edge, val)
}
func isEnded(levelsEnded []int, level int) bool {
for _, l := range levelsEnded {
if l == level {
return true
}
}
return false
}
type EdgeType string
var (
EdgeTypeLink EdgeType = "│"
EdgeTypeMid EdgeType = "├──"
EdgeTypeEnd EdgeType = "└──"
)
func New() Tree {
return &node{Value: "."}
}