fix(tree): print top node #443

pterm · Jan 22, 2023 · 591aebd · 591aebd
1 parent dbca1d8
commit 591aebd
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 32 deletions.
diff --git a/_examples/tree/demo/main.go b/_examples/tree/demo/main.go
@@ -2,39 +2,20 @@ package main
 
 import (
 	"github.com/pterm/pterm"
-	"github.com/pterm/pterm/putils"
 )
 
 func main() {
-	// You can use a LeveledList here, for easy generation.
-	leveledList := pterm.LeveledList{
-		pterm.LeveledListItem{Level: 0, Text: "C:"},
-		pterm.LeveledListItem{Level: 1, Text: "Users"},
-		pterm.LeveledListItem{Level: 1, Text: "Windows"},
-		pterm.LeveledListItem{Level: 1, Text: "Programs"},
-		pterm.LeveledListItem{Level: 1, Text: "Programs(x86)"},
-		pterm.LeveledListItem{Level: 1, Text: "dev"},
-		pterm.LeveledListItem{Level: 0, Text: "D:"},
-		pterm.LeveledListItem{Level: 0, Text: "E:"},
-		pterm.LeveledListItem{Level: 1, Text: "Movies"},
-		pterm.LeveledListItem{Level: 1, Text: "Music"},
-		pterm.LeveledListItem{Level: 2, Text: "LinkinPark"},
-		pterm.LeveledListItem{Level: 1, Text: "Games"},
-		pterm.LeveledListItem{Level: 2, Text: "Shooter"},
-		pterm.LeveledListItem{Level: 3, Text: "CallOfDuty"},
-		pterm.LeveledListItem{Level: 3, Text: "CS:GO"},
-		pterm.LeveledListItem{Level: 3, Text: "Battlefield"},
-		pterm.LeveledListItem{Level: 4, Text: "Battlefield 1"},
-		pterm.LeveledListItem{Level: 4, Text: "Battlefield 2"},
-		pterm.LeveledListItem{Level: 0, Text: "F:"},
-		pterm.LeveledListItem{Level: 1, Text: "dev"},
-		pterm.LeveledListItem{Level: 2, Text: "dops"},
-		pterm.LeveledListItem{Level: 2, Text: "PTerm"},
+	tree := pterm.TreeNode{
+		Text: "Top node",
+		Children: []pterm.TreeNode{{
+			Text: "Child node",
+			Children: []pterm.TreeNode{
+				{Text: "Grandchild node"},
+				{Text: "Grandchild node"},
+				{Text: "Grandchild node"},
+			},
+		}},
 	}
 
-	// Generate tree from LeveledList.
-	root := putils.TreeFromLeveledList(leveledList)
-
-	// Render TreePrinter
-	pterm.DefaultTree.WithRoot(root).Render()
+	pterm.DefaultTree.WithRoot(tree).Render()
 }
diff --git a/_examples/tree/from-leveled-list/main.go b/_examples/tree/from-leveled-list/main.go
@@ -0,0 +1,41 @@
+package main
+
+import (
+	"github.com/pterm/pterm"
+	"github.com/pterm/pterm/putils"
+)
+
+func main() {
+	// You can use a LeveledList here, for easy generation.
+	leveledList := pterm.LeveledList{
+		pterm.LeveledListItem{Level: 0, Text: "C:"},
+		pterm.LeveledListItem{Level: 1, Text: "Users"},
+		pterm.LeveledListItem{Level: 1, Text: "Windows"},
+		pterm.LeveledListItem{Level: 1, Text: "Programs"},
+		pterm.LeveledListItem{Level: 1, Text: "Programs(x86)"},
+		pterm.LeveledListItem{Level: 1, Text: "dev"},
+		pterm.LeveledListItem{Level: 0, Text: "D:"},
+		pterm.LeveledListItem{Level: 0, Text: "E:"},
+		pterm.LeveledListItem{Level: 1, Text: "Movies"},
+		pterm.LeveledListItem{Level: 1, Text: "Music"},
+		pterm.LeveledListItem{Level: 2, Text: "LinkinPark"},
+		pterm.LeveledListItem{Level: 1, Text: "Games"},
+		pterm.LeveledListItem{Level: 2, Text: "Shooter"},
+		pterm.LeveledListItem{Level: 3, Text: "CallOfDuty"},
+		pterm.LeveledListItem{Level: 3, Text: "CS:GO"},
+		pterm.LeveledListItem{Level: 3, Text: "Battlefield"},
+		pterm.LeveledListItem{Level: 4, Text: "Battlefield 1"},
+		pterm.LeveledListItem{Level: 4, Text: "Battlefield 2"},
+		pterm.LeveledListItem{Level: 0, Text: "F:"},
+		pterm.LeveledListItem{Level: 1, Text: "dev"},
+		pterm.LeveledListItem{Level: 2, Text: "dops"},
+		pterm.LeveledListItem{Level: 2, Text: "PTerm"},
+	}
+
+	// Generate tree from LeveledList.
+	root := putils.TreeFromLeveledList(leveledList)
+	root.Text = "Computer"
+
+	// Render TreePrinter
+	pterm.DefaultTree.WithRoot(root).Render()
+}
diff --git a/putils/tree_from_leveled_list.go b/putils/tree_from_leveled_list.go
@@ -10,7 +10,6 @@ func TreeFromLeveledList(leveledListItems pterm.LeveledList) pterm.TreeNode {
 
 	root := &pterm.TreeNode{
 		Children: []pterm.TreeNode{},
-		Text:     leveledListItems[0].Text,
 	}
 
 	for i, record := range leveledListItems {

diff --git a/tree_printer.go b/tree_printer.go
@@ -122,7 +122,12 @@ func (p TreePrinter) Srender() (string, error) {
 		p.TextStyle = NewStyle()
 	}
 
-	return walkOverTree(p.Root.Children, p, ""), nil
+	var result string
+	if p.Root.Text != "" {
+		result += p.Root.Text + "\n"
+	}
+	result += walkOverTree(p.Root.Children, p, "")
+	return result, nil
 }
 
 // walkOverTree is a recursive function,