diff --git a/tree/tree.go b/tree/tree.go new file mode 100644 index 00000000..b50ff6dd --- /dev/null +++ b/tree/tree.go @@ -0,0 +1,89 @@ +package tree + +import "strings" + +// Node is a node in a tree. +type Node interface { + String() string + Children() []Node +} + +// IndentFunc is the function that allow customization of the indentation of +// the tree. +type IndentFunc func(n Node, level, index int, last bool) string + +// StringNode implements the Node interface with String data. +type StringNode struct { + indentFunc IndentFunc + data *string + children []Node +} + +func (n StringNode) string(level, index int, last bool) string { + var s strings.Builder + + if n.data != nil { + if n.indentFunc != nil { + s.WriteString(n.indentFunc(n, level, index, last)) + } + s.WriteString(*n.data + "\n") + } + + for i, child := range n.children { + c := child.(*StringNode) + c.indentFunc = n.indentFunc + s.WriteString(c.string(level+1, i, i == len(n.children)-1)) + } + + return s.String() +} + +func (n StringNode) String() string { + return n.string(-1, 0, false) +} + +// Indent sets the indentation function for a string node / tree. +func (n *StringNode) Indent(indentFunc IndentFunc) *StringNode { + n.indentFunc = indentFunc + return n +} + +// Children returns the children of a string node. +func (n StringNode) Children() []Node { + return n.children +} + +func defaultIndentFunc(_ Node, level, _ int, last bool) string { + var s strings.Builder + + if level >= 0 { + s.WriteString(strings.Repeat("│ ", level)) + } + + if last { + s.WriteString("└── ") + } else { + s.WriteString("├── ") + } + + return s.String() +} + +// New returns a new tree. +func New(data ...any) *StringNode { + var children []Node + + for _, d := range data { + switch d := d.(type) { + case *StringNode: + children = append(children, d) + case string: + children = append(children, &StringNode{data: &d, indentFunc: defaultIndentFunc}) + } + } + + return &StringNode{ + children: children, + indentFunc: defaultIndentFunc, + } +} diff --git a/tree/tree_test.go b/tree/tree_test.go new file mode 100644 index 00000000..805a42b9 --- /dev/null +++ b/tree/tree_test.go @@ -0,0 +1,102 @@ +package tree + +import ( + "strings" + "testing" +) + +func TestTree(t *testing.T) { + tree := New( + "Foo", + "Bar", + New( + "Qux", + "Quux", + New( + "Foo", + "Bar", + ), + "Quuux", + ), + "Baz", + ) + + expected := strings.TrimPrefix(` +├── Foo +├── Bar +│ ├── Qux +│ ├── Quux +│ │ ├── Foo +│ │ └── Bar +│ └── Quuux +└── Baz +`, "\n") + + if tree.String() != expected { + t.Fatalf("expected:\n\n%s\n\ngot:\n\n%s\n", expected, tree) + } +} + +func TestTreeNil(t *testing.T) { + tree := New( + nil, + "Bar", + New( + "Qux", + "Quux", + New( + nil, + "Bar", + ), + "Quuux", + ), + "Baz", + ) + + expected := strings.TrimPrefix(` +├── Bar +│ ├── Qux +│ ├── Quux +│ │ └── Bar +│ └── Quuux +└── Baz +`, "\n") + + if tree.String() != expected { + t.Fatalf("expected:\n\n%s\n\ngot:\n\n%s\n", expected, tree) + } +} + +func TestTreeCustom(t *testing.T) { + tree := New( + "Foo", + "Bar", + New( + "Qux", + "Quux", + New( + "Foo", + "Bar", + ), + "Quuux", + ), + "Baz", + ).Indent(func(t Node, level, index int, last bool) string { + return strings.Repeat("-> ", level+1) + }) + + expected := strings.TrimPrefix(` +-> Foo +-> Bar +-> -> Qux +-> -> Quux +-> -> -> Foo +-> -> -> Bar +-> -> Quuux +-> Baz +`, "\n") + + if tree.String() != expected { + t.Fatalf("expected:\n\n%s\n\ngot:\n\n%s\n", expected, tree) + } +}