use crate::console::RenderContext;
use crate::renderable::{Renderable, Segment};
use crate::style::Style;
use crate::text::{Span, Text};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum GuideStyle {
Ascii,
#[default]
Unicode,
Bold,
Double,
}
impl GuideStyle {
fn chars(&self) -> TreeGuideChars {
match self {
GuideStyle::Ascii => TreeGuideChars {
vertical: '|',
horizontal: '-',
branch: '+',
last_branch: '\\',
space: ' ',
},
GuideStyle::Unicode => TreeGuideChars {
vertical: '│',
horizontal: '─',
branch: '├',
last_branch: '└',
space: ' ',
},
GuideStyle::Bold => TreeGuideChars {
vertical: '┃',
horizontal: '━',
branch: '┣',
last_branch: '┗',
space: ' ',
},
GuideStyle::Double => TreeGuideChars {
vertical: '║',
horizontal: '═',
branch: '╠',
last_branch: '╚',
space: ' ',
},
}
}
}
#[derive(Debug, Clone, Copy)]
struct TreeGuideChars {
vertical: char,
horizontal: char,
branch: char,
last_branch: char,
#[allow(dead_code)]
space: char,
}
#[derive(Debug, Clone)]
pub struct TreeNode {
label: Text,
children: Vec<TreeNode>,
style: Style,
expanded: bool,
}
impl TreeNode {
pub fn new<T: Into<Text>>(label: T) -> Self {
TreeNode {
label: label.into(),
children: Vec::new(),
style: Style::new(),
expanded: true,
}
}
pub fn add<T: Into<TreeNode>>(&mut self, child: T) -> &mut Self {
self.children.push(child.into());
self
}
pub fn with_child<T: Into<TreeNode>>(mut self, child: T) -> Self {
self.children.push(child.into());
self
}
pub fn with_children<I, T>(mut self, children: I) -> Self
where
I: IntoIterator<Item = T>,
T: Into<TreeNode>,
{
for child in children {
self.children.push(child.into());
}
self
}
pub fn style(mut self, style: Style) -> Self {
self.style = style;
self
}
pub fn expanded(mut self, expanded: bool) -> Self {
self.expanded = expanded;
self
}
pub fn has_children(&self) -> bool {
!self.children.is_empty()
}
}
impl<T: Into<Text>> From<T> for TreeNode {
fn from(label: T) -> Self {
TreeNode::new(label)
}
}
#[derive(Debug, Clone)]
pub struct Tree {
root: TreeNode,
guide_style: GuideStyle,
style: Style,
hide_root: bool,
}
impl Tree {
pub fn new<T: Into<TreeNode>>(root: T) -> Self {
Tree {
root: root.into(),
guide_style: GuideStyle::Unicode,
style: Style::new(),
hide_root: false,
}
}
pub fn guide_style(mut self, style: GuideStyle) -> Self {
self.guide_style = style;
self
}
pub fn style(mut self, style: Style) -> Self {
self.style = style;
self
}
pub fn hide_root(mut self, hide: bool) -> Self {
self.hide_root = hide;
self
}
pub fn add<T: Into<TreeNode>>(&mut self, child: T) -> &mut Self {
self.root.children.push(child.into());
self
}
fn render_node(
&self,
node: &TreeNode,
prefix: &str,
is_last: bool,
is_root: bool,
chars: &TreeGuideChars,
segments: &mut Vec<Segment>,
) {
if !is_root || !self.hide_root {
let mut spans = Vec::new();
if !is_root {
spans.push(Span::styled(prefix.to_string(), self.style));
let branch_char = if is_last {
chars.last_branch
} else {
chars.branch
};
spans.push(Span::styled(
format!("{}{}{} ", branch_char, chars.horizontal, chars.horizontal),
self.style,
));
}
for span in &node.label.spans {
let combined_style = node.style.combine(&span.style);
spans.push(Span::styled(span.text.to_string(), combined_style));
}
segments.push(Segment::line(spans));
}
if node.expanded {
let child_count = node.children.len();
for (i, child) in node.children.iter().enumerate() {
let is_last_child = i == child_count - 1;
let new_prefix = if is_root {
String::new()
} else {
let connector = if is_last {
" ".to_string()
} else {
format!("{} ", chars.vertical)
};
format!("{}{}", prefix, connector)
};
self.render_node(child, &new_prefix, is_last_child, false, chars, segments);
}
}
}
}
impl Renderable for Tree {
fn render(&self, _context: &RenderContext) -> Vec<Segment> {
let chars = self.guide_style.chars();
let mut segments = Vec::new();
self.render_node(&self.root, "", true, true, &chars, &mut segments);
segments
}
}
#[cfg(feature = "std")]
pub fn from_directory(path: &std::path::Path) -> std::io::Result<Tree> {
fn build_node(path: &std::path::Path) -> std::io::Result<TreeNode> {
let name = path
.file_name()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| path.to_string_lossy().to_string());
let mut node = TreeNode::new(name);
if path.is_dir() {
let mut entries: Vec<_> = std::fs::read_dir(path)?.filter_map(|e| e.ok()).collect();
entries.sort_by_key(|e| e.file_name());
for entry in entries {
let child = build_node(&entry.path())?;
node.children.push(child);
}
}
Ok(node)
}
let root = build_node(path)?;
Ok(Tree::new(root))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tree_simple() {
let tree = Tree::new("root").guide_style(GuideStyle::Unicode);
let context = RenderContext {
width: 40,
height: None,
};
let segments = tree.render(&context);
assert_eq!(segments.len(), 1);
assert!(segments[0].plain_text().contains("root"));
}
#[test]
fn test_tree_with_children() {
let mut tree = Tree::new("root");
tree.add(TreeNode::new("child1"));
tree.add(TreeNode::new("child2"));
let context = RenderContext {
width: 40,
height: None,
};
let segments = tree.render(&context);
assert_eq!(segments.len(), 3);
}
#[test]
fn test_tree_nested() {
let child1 = TreeNode::new("child1")
.with_child("grandchild1")
.with_child("grandchild2");
let tree = Tree::new(TreeNode::new("root").with_child(child1));
let context = RenderContext {
width: 40,
height: None,
};
let segments = tree.render(&context);
assert_eq!(segments.len(), 4);
let text: String = segments.iter().map(|s| s.plain_text()).collect();
assert!(text.contains("├"));
assert!(text.contains("└"));
}
#[test]
fn test_tree_hide_root() {
let mut tree = Tree::new("root").hide_root(true);
tree.add("child1");
tree.add("child2");
let context = RenderContext {
width: 40,
height: None,
};
let segments = tree.render(&context);
let text: String = segments.iter().map(|s| s.plain_text()).collect();
assert!(!text.contains("root"));
assert!(text.contains("child1"));
}
}