From 7357ed464986a9833c6717ca453572719ad1628e Mon Sep 17 00:00:00 2001
From: Igor German <igor.german@synesis.ru>
Date: Thu, 2 May 2019 15:50:48 +0300
Subject: [PATCH] resolve pr comments

---
 _examples/tree.go |  2 +-
 theme.go          | 12 ++++++------
 widgets/tree.go   | 32 ++++++++++++++++++++------------
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/_examples/tree.go b/_examples/tree.go
index 18753c9..b572925 100644
--- a/_examples/tree.go
+++ b/_examples/tree.go
@@ -70,7 +70,7 @@ func main() {
 	l := widgets.NewTree()
 	l.TextStyle = ui.NewStyle(ui.ColorYellow)
 	l.WrapText = false
-	l.Nodes = nodes
+	l.SetNodes(nodes)
 
 	x, y := ui.TerminalDimensions()
 
diff --git a/theme.go b/theme.go
index 8d75174..60745f4 100644
--- a/theme.go
+++ b/theme.go
@@ -68,9 +68,9 @@ type ListTheme struct {
 }
 
 type TreeTheme struct {
-	Text     Style
-	Collapse rune
-	Expand   rune
+	Text      Style
+	Collapsed rune
+	Expanded  rune
 }
 
 type ParagraphTheme struct {
@@ -130,9 +130,9 @@ var Theme = RootTheme{
 	},
 
 	Tree: TreeTheme{
-		Text:     NewStyle(ColorWhite),
-		Collapse: COLLAPSED,
-		Expand:   EXPANDED,
+		Text:      NewStyle(ColorWhite),
+		Collapsed: COLLAPSED,
+		Expanded:  EXPANDED,
 	},
 
 	StackedBarChart: StackedBarChartTheme{
diff --git a/widgets/tree.go b/widgets/tree.go
index 6217a68..3060e34 100644
--- a/widgets/tree.go
+++ b/widgets/tree.go
@@ -17,7 +17,7 @@ type TreeNode struct {
 	Expanded bool
 	Nodes    []*TreeNode
 
-	//level stores the node level in the tree.
+	// level stores the node level in the tree.
 	level int
 }
 
@@ -32,9 +32,9 @@ func (self *TreeNode) parseStyles(style Style) []Cell {
 	} else {
 		sb.WriteString(strings.Repeat(treeIndent, self.level))
 		if self.Expanded {
-			sb.WriteRune(Theme.Tree.Expand)
+			sb.WriteRune(Theme.Tree.Expanded)
 		} else {
-			sb.WriteRune(Theme.Tree.Collapse)
+			sb.WriteRune(Theme.Tree.Collapsed)
 		}
 		sb.WriteByte(' ')
 	}
@@ -45,15 +45,15 @@ func (self *TreeNode) parseStyles(style Style) []Cell {
 // Tree is a tree widget.
 type Tree struct {
 	Block
-	Nodes            []*TreeNode
 	TextStyle        Style
 	SelectedRowStyle Style
 	WrapText         bool
 	SelectedRow      int
-	topRow           int
 
-	//rows is flatten nodes for rendering.
-	rows []*TreeNode
+	nodes []*TreeNode
+	// rows is flatten nodes for rendering.
+	rows   []*TreeNode
+	topRow int
 }
 
 // NewTree creates a new Tree widget.
@@ -66,9 +66,14 @@ func NewTree() *Tree {
 	}
 }
 
+func (self *Tree) SetNodes(nodes []*TreeNode) {
+	self.nodes = nodes
+	self.prepareNodes()
+}
+
 func (self *Tree) prepareNodes() {
-	self.rows = self.rows[:0]
-	for _, node := range self.Nodes {
+	self.rows = make([]*TreeNode, 0)
+	for _, node := range self.nodes {
 		self.prepareNode(node, 0)
 	}
 }
@@ -85,7 +90,7 @@ func (self *Tree) prepareNode(node *TreeNode, level int) {
 }
 
 func (self *Tree) Walk(fn TreeWalkFn) {
-	for _, n := range self.Nodes {
+	for _, n := range self.nodes {
 		if !self.walk(n, fn) {
 			break
 		}
@@ -108,8 +113,6 @@ func (self *Tree) walk(n *TreeNode, fn TreeWalkFn) bool {
 
 func (self *Tree) Draw(buf *Buffer) {
 	self.Block.Draw(buf)
-	self.prepareNodes()
-
 	point := self.Inner.Min
 
 	// adjusts view into widget
@@ -209,6 +212,7 @@ func (self *Tree) ScrollBottom() {
 
 func (self *Tree) Collapse() {
 	self.rows[self.SelectedRow].Expanded = false
+	self.prepareNodes()
 }
 
 func (self *Tree) Expand() {
@@ -216,6 +220,7 @@ func (self *Tree) Expand() {
 	if len(node.Nodes) > 0 {
 		self.rows[self.SelectedRow].Expanded = true
 	}
+	self.prepareNodes()
 }
 
 func (self *Tree) ToggleExpand() {
@@ -223,6 +228,7 @@ func (self *Tree) ToggleExpand() {
 	if len(node.Nodes) > 0 {
 		node.Expanded = !node.Expanded
 	}
+	self.prepareNodes()
 }
 
 func (self *Tree) ExpandAll() {
@@ -232,6 +238,7 @@ func (self *Tree) ExpandAll() {
 		}
 		return true
 	})
+	self.prepareNodes()
 }
 
 func (self *Tree) CollapseAll() {
@@ -239,4 +246,5 @@ func (self *Tree) CollapseAll() {
 		n.Expanded = false
 		return true
 	})
+	self.prepareNodes()
 }