33package tree
44
55import (
6+ "fmt"
67 "strings"
78
89 tea "charm.land/bubbletea/v2"
910 "charm.land/lipgloss/v2"
1011 ltree "charm.land/lipgloss/v2/tree"
12+ "github.com/charmbracelet/x/ansi"
1113
1214 "charm.land/bubbles/v2/help"
1315 "charm.land/bubbles/v2/key"
@@ -245,6 +247,7 @@ func (m *Model) SetCursorCharacter(character string) {
245247// SetNodes sets the tree to the given root node.
246248func (m * Model ) SetNodes (t * Node ) {
247249 m .root = t
250+ m .root .value = t .Value ()
248251 if m .enumerator != nil {
249252 m .root .Enumerator (* m .enumerator )
250253 }
@@ -254,6 +257,7 @@ func (m *Model) SetNodes(t *Node) {
254257 m .setAttributes ()
255258 m .updateStyles ()
256259 m .updateViewport (0 )
260+ m .setRootStyles (m .styles )
257261}
258262
259263// SetAdditionalFullHelpKeys sets additional key mappings for the full help view.
@@ -366,7 +370,7 @@ func (m *Model) updateViewport(movement int) {
366370 lipgloss .JoinHorizontal (
367371 lipgloss .Top ,
368372 cursor ,
369- m .styles .TreeStyle .Render (m .root .String ()),
373+ m .styles .TreeStyle .Width ( m . width ). MaxWidth ( m . width ). Render (m .root .String ()),
370374 ),
371375 )
372376
@@ -418,20 +422,51 @@ func (m *Model) SetStyles(styles Styles) {
418422 return styles .RootNodeStyle
419423 }
420424 }
425+ if styles .EnumeratorStyleFunc != nil {
426+ styles .enumeratorFunc = styles .EnumeratorStyleFunc
427+ } else {
428+ styles .enumeratorFunc = func (_ Nodes , _ int ) lipgloss.Style {
429+ return styles .EnumeratorStyle
430+ }
431+ }
432+ if styles .SelectedEnumeratorStyleFunc != nil {
433+ styles .selectedEnumeratorFunc = styles .SelectedEnumeratorStyleFunc
434+ } else {
435+ styles .selectedEnumeratorFunc = func (_ Nodes , _ int ) lipgloss.Style {
436+ return styles .SelectedEnumeratorStyle
437+ }
438+ }
439+ if styles .IndenterStyleFunc != nil {
440+ styles .indenterFunc = styles .IndenterStyleFunc
441+ } else {
442+ styles .indenterFunc = func (_ Nodes , _ int ) lipgloss.Style {
443+ return styles .IndenterStyle
444+ }
445+ }
421446
447+ m .setRootStyles (styles )
448+ m .styles = styles
449+
450+ // call SetSize as it takes into account width/height of the styles frame sizes
451+ m .SetSize (m .width , m .height )
452+ m .updateViewport (0 )
453+ }
454+
455+ func (m * Model ) setRootStyles (styles Styles ) {
422456 if m .root != nil {
423- m .root .EnumeratorStyle (styles .EnumeratorStyle )
424- m .root .IndenterStyle (styles .IndenterStyle )
457+ m .root .EnumeratorStyleFunc (func (children Nodes , i int ) lipgloss.Style {
458+ child := children .At (i )
459+ return child .getEnumeratorStyle ()
460+ })
461+ m .root .IndenterStyleFunc (func (children Nodes , i int ) lipgloss.Style {
462+ child := children .At (i )
463+ return child .getIndenterStyle ()
464+ })
425465 m .root .ItemStyleFunc (func (children Nodes , i int ) lipgloss.Style {
426466 child := children .At (i )
427467 return child .getStyle ()
428468 })
429469 }
430-
431- m .styles = styles
432- // call SetSize as it takes into account width/height of the styles frame sizes
433- m .SetSize (m .width , m .height )
434- m .updateViewport (0 )
435470}
436471
437472// SetShowHelp shows or hides the help view.
@@ -464,8 +499,10 @@ func (m *Model) SetHeight(height int) {
464499func (m * Model ) SetSize (width , height int ) {
465500 m .width = width
466501 m .height = height
467- m .root .tree .Width (width - lipgloss .Width (m .cursorView ()) -
468- m .styles .TreeStyle .GetHorizontalFrameSize ())
502+ if m .root != nil {
503+ m .root .tree .Width (width - lipgloss .Width (m .cursorView ()) -
504+ m .styles .TreeStyle .GetHorizontalFrameSize ())
505+ }
469506
470507 m .viewport .SetWidth (width )
471508 hv := 0
@@ -589,6 +626,11 @@ func (m *Model) YOffset() int {
589626 return m .yOffset
590627}
591628
629+ func (m * Model ) SetYOffset (yoffset int ) {
630+ movement := yoffset - m .yOffset
631+ m .updateViewport (movement )
632+ }
633+
592634// Node returns the item at the given yoffset.
593635func (m * Model ) Node (yoffset int ) * Node {
594636 return findNode (m .root , yoffset )
@@ -602,26 +644,42 @@ func (m *Model) NodeAtCurrentOffset() *Node {
602644// Enumerator sets the enumerator for the tree.
603645func (m * Model ) Enumerator (enumerator ltree.Enumerator ) * Model {
604646 m .enumerator = & enumerator
605- m .root .Enumerator (enumerator )
647+ if m .root != nil {
648+ m .root .Enumerator (enumerator )
649+ }
606650 return m
607651}
608652
609653// Indenter sets the indenter for the tree.
610654func (m * Model ) Indenter (indenter ltree.Indenter ) * Model {
611655 m .indenter = & indenter
612- m .root .Indenter (indenter )
656+ if m .root != nil {
657+ m .root .Indenter (indenter )
658+ }
613659 return m
614660}
615661
616662// Since the selected node changes, we need to capture m.yOffset in the
617663// style function's closure again.
618664func (m * Model ) updateStyles () {
665+ opts := m .getItemOpts ()
619666 if m .root != nil {
620- m .root .RootStyle (m .rootStyle ())
667+ m .root .opts = * opts
668+ rs := m .rootStyle ()
669+ m .root .RootStyle (rs )
670+
671+ root := ""
672+ switch val := m .root .value .(type ) {
673+ case fmt.Stringer :
674+ root = rs .Render (val .String ())
675+ case string :
676+ root = rs .Render (ansi .Strip (val ))
677+ }
678+ root = lipgloss .JoinHorizontal (lipgloss .Left , m .root .Indicator (), root )
679+ m .root .tree .SetValue (root )
621680 }
622681
623682 items := m .AllNodes ()
624- opts := m .getItemOpts ()
625683 for _ , item := range items {
626684 item .opts = * opts
627685 }
@@ -637,11 +695,17 @@ func (m *Model) getItemOpts() *itemOptions {
637695}
638696
639697func (m * Model ) rootStyle () lipgloss.Style {
640- if m .root .yOffset == m .yOffset {
641- return m .styles .selectedNodeFunc (Nodes {m .root }, 0 )
698+ nodes := Nodes {m .root }
699+ if m .root .yOffset == m .yOffset && m .styles .selectedNodeFunc != nil {
700+ return m .styles .rootNodeFunc (nodes , 0 ).Inherit (
701+ m .styles .selectedNodeFunc (Nodes {m .root }, 0 ))
702+ }
703+
704+ if m .styles .rootNodeFunc != nil {
705+ return m .styles .rootNodeFunc (Nodes {m .root }, 0 )
642706 }
643707
644- return m . styles . rootNodeFunc ( Nodes { m . root }, 0 )
708+ return lipgloss . NewStyle ( )
645709}
646710
647711// findNode starts a DFS search for the node with the given yOffset
0 commit comments