use crate::console::{Console, ConsoleOptions};
use crate::renderables::Renderable;
use crate::segment::Segment;
use crate::style::Style;
use crate::text::Text;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TreeGuides {
Ascii,
#[default]
Unicode,
Bold,
Double,
Rounded,
}
impl TreeGuides {
#[must_use]
pub const fn vertical(&self) -> &str {
match self {
Self::Ascii => "| ",
Self::Unicode | Self::Rounded => "\u{2502} ", Self::Bold => "\u{2503} ", Self::Double => "\u{2551} ", }
}
#[must_use]
pub const fn branch(&self) -> &str {
match self {
Self::Ascii => "+-- ",
Self::Unicode => "\u{251C}\u{2500}\u{2500} ", Self::Bold => "\u{2523}\u{2501}\u{2501} ", Self::Double => "\u{2560}\u{2550}\u{2550} ", Self::Rounded => "\u{251C}\u{2500}\u{2500} ", }
}
#[must_use]
pub const fn last(&self) -> &str {
match self {
Self::Ascii => "`-- ",
Self::Unicode => "\u{2514}\u{2500}\u{2500} ", Self::Bold => "\u{2517}\u{2501}\u{2501} ", Self::Double => "\u{255A}\u{2550}\u{2550} ", Self::Rounded => "\u{2570}\u{2500}\u{2500} ", }
}
#[must_use]
pub const fn space(&self) -> &'static str {
" "
}
}
#[derive(Debug, Clone)]
pub struct TreeNode {
label: Text,
children: Vec<TreeNode>,
expanded: bool,
icon: Option<String>,
icon_style: Style,
}
impl TreeNode {
#[must_use]
pub fn new(label: impl Into<Text>) -> Self {
Self {
label: label.into(),
children: Vec::new(),
expanded: true,
icon: None,
icon_style: Style::new(),
}
}
#[must_use]
pub fn with_icon(icon: impl Into<String>, label: impl Into<Text>) -> Self {
Self {
label: label.into(),
children: Vec::new(),
expanded: true,
icon: Some(icon.into()),
icon_style: Style::new(),
}
}
#[must_use]
pub fn child(mut self, node: TreeNode) -> Self {
self.children.push(node);
self
}
#[must_use]
pub fn children(mut self, nodes: impl IntoIterator<Item = TreeNode>) -> Self {
self.children.extend(nodes);
self
}
#[must_use]
pub fn icon(mut self, icon: impl Into<String>) -> Self {
self.icon = Some(icon.into());
self
}
#[must_use]
pub fn icon_style(mut self, style: Style) -> Self {
self.icon_style = style;
self
}
#[must_use]
pub fn expanded(mut self, expanded: bool) -> Self {
self.expanded = expanded;
self
}
#[must_use]
pub fn collapsed(self) -> Self {
self.expanded(false)
}
#[must_use]
pub fn label(&self) -> &Text {
&self.label
}
#[must_use]
pub fn children_nodes(&self) -> &[TreeNode] {
&self.children
}
#[must_use]
pub fn has_children(&self) -> bool {
!self.children.is_empty()
}
#[must_use]
pub fn is_expanded(&self) -> bool {
self.expanded
}
#[must_use]
pub fn get_icon(&self) -> Option<&str> {
self.icon.as_deref()
}
}
#[derive(Debug, Clone)]
pub struct Tree {
root: TreeNode,
guides: TreeGuides,
guide_style: Style,
show_root: bool,
highlight_style: Option<Style>,
max_depth: isize,
}
impl Default for Tree {
fn default() -> Self {
Self {
root: TreeNode::new("root"),
guides: TreeGuides::default(),
guide_style: Style::new(),
show_root: true,
highlight_style: None,
max_depth: -1,
}
}
}
impl Tree {
#[must_use]
pub fn new(root: TreeNode) -> Self {
Self {
root,
..Self::default()
}
}
#[must_use]
pub fn with_label(label: impl Into<Text>) -> Self {
Self::new(TreeNode::new(label))
}
#[must_use]
pub fn guides(mut self, guides: TreeGuides) -> Self {
self.guides = guides;
self
}
#[must_use]
pub fn guide_style(mut self, style: Style) -> Self {
self.guide_style = style;
self
}
#[must_use]
pub fn show_root(mut self, show: bool) -> Self {
self.show_root = show;
self
}
#[must_use]
pub fn hide_root(self) -> Self {
self.show_root(false)
}
#[must_use]
pub fn highlight_style(mut self, style: Style) -> Self {
self.highlight_style = Some(style);
self
}
#[must_use]
pub fn max_depth(mut self, depth: isize) -> Self {
self.max_depth = depth;
self
}
#[must_use]
pub fn child(mut self, node: TreeNode) -> Self {
self.root.children.push(node);
self
}
#[must_use]
pub fn children(mut self, nodes: impl IntoIterator<Item = TreeNode>) -> Self {
self.root.children.extend(nodes);
self
}
#[must_use]
pub fn render(&self) -> Vec<Segment<'_>> {
let mut segments = Vec::new();
let prefix_stack: Vec<bool> = Vec::new();
if self.show_root {
self.render_node(&self.root, &mut segments, &prefix_stack, true, 0);
} else {
let children = &self.root.children;
for (i, child) in children.iter().enumerate() {
let is_last = i == children.len() - 1;
self.render_node(child, &mut segments, &prefix_stack, is_last, 0);
}
}
segments
}
fn sanitize_label(label: &Text) -> Text {
if !label.plain().contains('\n') {
return label.clone();
}
let mut sanitized = Text::new(label.plain().replace('\n', " "));
sanitized.set_style(label.style().clone());
sanitized.justify = label.justify;
sanitized.overflow = label.overflow;
sanitized.no_wrap = label.no_wrap;
sanitized.end.clone_from(&label.end);
sanitized.tab_size = label.tab_size;
for span in label.spans() {
sanitized.stylize(span.start, span.end, span.style.clone());
}
sanitized
}
#[expect(
clippy::cast_possible_wrap,
reason = "tree depth will never exceed isize::MAX"
)]
fn render_node<'a>(
&'a self,
node: &'a TreeNode,
segments: &mut Vec<Segment<'a>>,
prefix_stack: &[bool],
is_last: bool,
depth: usize,
) {
if self.max_depth >= 0 && depth as isize > self.max_depth {
return;
}
for &has_more_siblings in prefix_stack {
let guide = if has_more_siblings {
self.guides.vertical()
} else {
self.guides.space()
};
segments.push(Segment::new(guide, Some(self.guide_style.clone())));
}
if depth > 0 || !self.show_root {
let guide = if is_last {
self.guides.last()
} else {
self.guides.branch()
};
segments.push(Segment::new(guide, Some(self.guide_style.clone())));
}
if let Some(icon) = node.get_icon() {
segments.push(Segment::new(
format!("{icon} "),
Some(node.icon_style.clone()),
));
}
let label_text = Self::sanitize_label(&node.label);
let mut label_segments: Vec<Segment<'static>> = label_text
.render("")
.into_iter()
.map(Segment::into_owned)
.collect();
if let Some(ref highlight) = self.highlight_style {
for segment in &mut label_segments {
if !segment.is_control() {
segment.style = Some(match segment.style.take() {
Some(existing) => existing.combine(highlight),
None => highlight.clone(),
});
}
}
}
for segment in label_segments {
segments.push(segment);
}
if node.has_children() && !node.is_expanded() {
segments.push(Segment::new(" [...]", Some(self.guide_style.clone())));
}
segments.push(Segment::line());
if node.is_expanded() {
let children = &node.children;
let mut new_prefix_stack = prefix_stack.to_vec();
if !(self.show_root && depth == 0) {
new_prefix_stack.push(!is_last);
}
for (i, child) in children.iter().enumerate() {
let child_is_last = i == children.len() - 1;
self.render_node(child, segments, &new_prefix_stack, child_is_last, depth + 1);
}
}
}
#[must_use]
pub fn render_plain(&self) -> String {
self.render()
.into_iter()
.map(|seg| seg.text.into_owned())
.collect()
}
}
impl Renderable for Tree {
fn render<'a>(&'a self, _console: &Console, _options: &ConsoleOptions) -> Vec<Segment<'a>> {
self.render()
}
}
#[must_use]
pub fn file_tree(root: &str, entries: &[(&str, bool)]) -> Tree {
let mut root_node = TreeNode::with_icon("📁", root);
for (path, is_dir) in entries {
let icon = if *is_dir { "📁" } else { "📄" };
root_node = root_node.child(TreeNode::with_icon(icon, *path));
}
Tree::new(root_node)
}
#[must_use]
pub fn ascii_tree(root: TreeNode) -> Tree {
Tree::new(root).guides(TreeGuides::Ascii)
}
#[must_use]
pub fn rounded_tree(root: TreeNode) -> Tree {
Tree::new(root).guides(TreeGuides::Rounded)
}
#[must_use]
pub fn bold_tree(root: TreeNode) -> Tree {
Tree::new(root).guides(TreeGuides::Bold)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tree_node_new() {
let node = TreeNode::new("test");
assert_eq!(node.label().plain(), "test");
assert!(node.children_nodes().is_empty());
assert!(node.is_expanded());
}
#[test]
fn test_tree_node_with_icon() {
let node = TreeNode::with_icon("📁", "folder");
assert_eq!(node.label().plain(), "folder");
assert_eq!(node.get_icon(), Some("📁"));
}
#[test]
fn test_tree_node_children() {
let node = TreeNode::new("root")
.child(TreeNode::new("child1"))
.child(TreeNode::new("child2"));
assert_eq!(node.children_nodes().len(), 2);
assert!(node.has_children());
}
#[test]
fn test_tree_node_collapsed() {
let node = TreeNode::new("test").collapsed();
assert!(!node.is_expanded());
}
#[test]
fn test_tree_new() {
let tree = Tree::with_label("root");
assert!(tree.show_root);
assert_eq!(tree.guides, TreeGuides::Unicode);
}
#[test]
fn test_tree_guides_ascii() {
let guides = TreeGuides::Ascii;
assert_eq!(guides.vertical(), "| ");
assert_eq!(guides.branch(), "+-- ");
assert_eq!(guides.last(), "`-- ");
assert_eq!(guides.space(), " ");
}
#[test]
fn test_tree_guides_unicode() {
let guides = TreeGuides::Unicode;
assert!(guides.vertical().starts_with('\u{2502}')); assert!(guides.branch().starts_with('\u{251C}')); assert!(guides.last().starts_with('\u{2514}')); }
#[test]
fn test_tree_render_simple() {
let tree = Tree::with_label("root")
.child(TreeNode::new("child1"))
.child(TreeNode::new("child2"));
let segments = tree.render();
assert!(!segments.is_empty());
let plain = tree.render_plain();
assert!(plain.contains("root"));
assert!(plain.contains("child1"));
assert!(plain.contains("child2"));
}
#[test]
fn test_tree_render_preserves_spans() {
use crate::style::Attributes;
let mut label = Text::new("root");
label.stylize(0, 4, Style::new().bold());
let tree = Tree::new(TreeNode::new(label));
let segments = tree.render();
let has_bold = segments.iter().any(|seg| {
seg.text.contains("root")
&& seg
.style
.as_ref()
.is_some_and(|style| style.attributes.contains(Attributes::BOLD))
});
assert!(has_bold);
}
#[test]
fn test_tree_render_preserves_spans_after_newline_sanitization() {
use crate::style::Attributes;
let mut label = Text::new("root\nnode");
label.stylize_all(Style::new().bold());
label.stylize(5, 9, Style::new().italic());
let tree = Tree::new(TreeNode::new(label));
let rendered = tree.render_plain();
assert!(rendered.contains("root node"));
assert!(!rendered.contains("root\nnode"));
let segments = tree.render();
let has_italic_node = segments.iter().any(|seg| {
seg.text.contains("node")
&& seg
.style
.as_ref()
.is_some_and(|style| style.attributes.contains(Attributes::ITALIC))
});
assert!(has_italic_node);
}
#[test]
fn test_tree_render_nested() {
let tree =
Tree::with_label("root").child(TreeNode::new("parent").child(TreeNode::new("child")));
let plain = tree.render_plain();
assert!(plain.contains("root"));
assert!(plain.contains("parent"));
assert!(plain.contains("child"));
}
#[test]
fn test_tree_hide_root() {
let tree = Tree::with_label("root")
.hide_root()
.child(TreeNode::new("visible"));
let plain = tree.render_plain();
assert!(!plain.contains("root"));
assert!(plain.contains("visible"));
}
#[test]
fn test_tree_collapsed_node() {
let tree = Tree::with_label("root").child(
TreeNode::new("collapsed")
.collapsed()
.child(TreeNode::new("hidden")),
);
let plain = tree.render_plain();
assert!(plain.contains("collapsed"));
assert!(plain.contains("[...]"));
assert!(!plain.contains("hidden"));
}
#[test]
fn test_tree_max_depth() {
let tree = Tree::with_label("root")
.max_depth(1)
.child(TreeNode::new("level1").child(TreeNode::new("level2")));
let plain = tree.render_plain();
assert!(plain.contains("root"));
assert!(plain.contains("level1"));
assert!(!plain.contains("level2"));
}
#[test]
fn test_tree_ascii_style() {
let tree = ascii_tree(TreeNode::new("root").child(TreeNode::new("child")));
let plain = tree.render_plain();
assert!(plain.contains("+--") || plain.contains("`--"));
}
#[test]
fn test_tree_with_icons() {
let tree = Tree::with_label("project")
.child(TreeNode::with_icon("📁", "src"))
.child(TreeNode::with_icon("📄", "README.md"));
let plain = tree.render_plain();
assert!(plain.contains("📁"));
assert!(plain.contains("📄"));
assert!(plain.contains("src"));
assert!(plain.contains("README.md"));
}
#[test]
fn test_file_tree() {
let tree = file_tree("project", &[("src", true), ("Cargo.toml", false)]);
let plain = tree.render_plain();
assert!(plain.contains("project"));
assert!(plain.contains("src"));
assert!(plain.contains("Cargo.toml"));
}
#[test]
fn test_tree_complex_structure() {
let tree = Tree::with_label("root")
.child(
TreeNode::new("branch1")
.child(TreeNode::new("leaf1"))
.child(TreeNode::new("leaf2")),
)
.child(
TreeNode::new("branch2")
.child(TreeNode::new("sub-branch").child(TreeNode::new("deep-leaf"))),
)
.child(TreeNode::new("leaf3"));
let plain = tree.render_plain();
assert!(plain.contains("root"));
assert!(plain.contains("branch1"));
assert!(plain.contains("branch2"));
assert!(plain.contains("leaf1"));
assert!(plain.contains("leaf2"));
assert!(plain.contains("leaf3"));
assert!(plain.contains("sub-branch"));
assert!(plain.contains("deep-leaf"));
}
#[test]
fn test_tree_empty_root() {
let tree = Tree::with_label("");
let plain = tree.render_plain();
let _ = plain;
}
#[test]
fn test_tree_single_node() {
let tree = Tree::with_label("single");
let plain = tree.render_plain();
assert!(plain.contains("single"));
assert!(!plain.contains("├──"));
assert!(!plain.contains("└──"));
}
#[test]
fn test_tree_wide_unicode_labels() {
let tree = Tree::with_label("项目") .child(TreeNode::new("源代码")) .child(TreeNode::new("文档"));
let plain = tree.render_plain();
assert!(plain.contains("项目"));
assert!(plain.contains("源代码"));
assert!(plain.contains("文档"));
}
#[test]
fn test_tree_emoji_labels() {
let tree = Tree::with_label("📁 Root")
.child(TreeNode::new("📄 File"))
.child(TreeNode::new("🔧 Config"));
let plain = tree.render_plain();
assert!(plain.contains("📁"));
assert!(plain.contains("📄"));
assert!(plain.contains("🔧"));
}
#[test]
fn test_tree_guides_bold() {
let guides = TreeGuides::Bold;
assert_eq!(guides.vertical(), "┃ ");
assert_eq!(guides.branch(), "┣━━ ");
assert_eq!(guides.last(), "┗━━ ");
assert_eq!(guides.space(), " ");
}
#[test]
fn test_tree_guides_double() {
let guides = TreeGuides::Double;
assert_eq!(guides.vertical(), "║ ");
assert_eq!(guides.branch(), "╠══ ");
assert_eq!(guides.last(), "╚══ ");
assert_eq!(guides.space(), " ");
}
#[test]
fn test_tree_guides_rounded() {
let guides = TreeGuides::Rounded;
assert_eq!(guides.vertical(), "│ ");
assert_eq!(guides.branch(), "├── ");
assert_eq!(guides.last(), "╰── "); assert_eq!(guides.space(), " ");
}
}