filetree/
node.rs

1use core::fmt;
2
3#[cfg(feature = "color")]
4use colored::{ColoredString, Colorize};
5
6use serde_yaml::Value;
7use uuid::Uuid;
8
9use crate::{errors::Error, filetree::Filetree};
10
11const TAB_NODE_BELOW: &str = "│   ";
12const TAB_NO_NODE_BELOW: &str = "    ";
13const NODE_ARM_LAST: &str = "└── ";
14const NODE_ARM: &str = "├── ";
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17struct NodeName(String);
18
19#[cfg(feature = "color")]
20impl NodeName {
21    fn colored(&self, is_directory: bool) -> ColoredString {
22        if is_directory {
23            self.0.cyan().bold()
24        } else {
25            self.0.clone().into()
26        }
27    }
28}
29
30impl From<NodeName> for String {
31    fn from(value: NodeName) -> Self {
32        value.0
33    }
34}
35
36impl fmt::Display for NodeName {
37    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38        write!(f, "{}", self.0)
39    }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct Node {
44    uuid: Uuid,
45    name: NodeName,
46    children: Vec<Node>,
47    is_directory: bool,
48}
49
50impl Node {
51    pub fn new(name: String, children: Vec<Node>, is_directory: bool) -> Node {
52        Node {
53            uuid: Uuid::new_v4(),
54            name: NodeName(name),
55            children,
56            is_directory,
57        }
58    }
59
60    pub fn parent<'a>(&'a self, filetree: &'a Filetree) -> Option<&Node> {
61        filetree.root().get_parent_of(self)
62    }
63
64    pub(crate) fn add_children(&mut self, children: &mut Vec<Node>) {
65        self.children.append(children);
66    }
67
68    pub fn has_child(&self, child: &Node) -> bool {
69        self.children.contains(child)
70    }
71
72    pub fn get_parent_of(&self, child: &Node) -> Option<&Node> {
73        if self.has_child(child) {
74            Some(self)
75        } else {
76            for node in self.children.iter() {
77                if let Some(parent) = node.get_parent_of(child) {
78                    return Some(parent);
79                }
80            }
81            None
82        }
83    }
84
85    fn next_sibling<'a>(&'a self, filetree: &'a Filetree) -> Option<&Node> {
86        let parent = self.parent(filetree)?;
87        let self_position = parent.children.iter().position(|child| child == self)?;
88        parent.children.get(self_position + 1)
89    }
90
91    fn next_pibling<'a>(&'a self, filetree: &'a Filetree) -> Option<&Node> {
92        self.parent(filetree)?.next_sibling(filetree)
93    }
94
95    fn modify_parts_from_parent<'a, T>(
96        &self,
97        parent_parts: &mut Vec<T>,
98        parent: &Node,
99        filetree: &Filetree,
100    ) where
101        T: From<&'a str>,
102    {
103        // Use the parts from parent, except for the last 2
104        if parent_parts.len() >= 2 {
105            parent_parts.drain(parent_parts.len() - 2..);
106        } else if !parent_parts.is_empty() {
107            parent_parts.drain(parent_parts.len() - 1..);
108        }
109
110        if self.next_pibling(filetree).is_some() {
111            parent_parts.push(TAB_NODE_BELOW.into());
112        } else if parent.parent(filetree).is_some() {
113            // if it has a parent and the parent is not the root
114            parent_parts.push(TAB_NO_NODE_BELOW.into());
115        }
116
117        if self.next_sibling(filetree).is_some() {
118            parent_parts.push(NODE_ARM.into());
119        } else {
120            parent_parts.push(NODE_ARM_LAST.into());
121        }
122    }
123
124    fn display_parts<'a>(&'a self, filetree: &'a Filetree) -> Vec<String> {
125        let Some(parent) = self.parent(filetree) else {
126            return vec![self.name.to_string()];
127        };
128        let mut parts = parent.display_parts(filetree);
129        self.modify_parts_from_parent(&mut parts, parent, filetree);
130        parts.push(self.name.to_string());
131        parts
132    }
133
134    #[cfg(feature = "color")]
135    fn display_parts_color<'a>(&'a self, filetree: &'a Filetree) -> Vec<ColoredString> {
136        let Some(parent) = self.parent(filetree) else {
137            return vec![self.name.colored(self.is_directory)];
138        };
139        let mut parts = parent.display_parts_color(filetree);
140        self.modify_parts_from_parent(&mut parts, parent, filetree);
141        parts.push(self.name.colored(self.is_directory));
142        parts
143    }
144
145    pub(crate) fn nodes(&self) -> Vec<&Node> {
146        let mut nodes = vec![self];
147        for child in self.children.iter().flat_map(|child| child.nodes()) {
148            nodes.push(child);
149        }
150        nodes
151    }
152
153    pub(crate) fn treeline(&self, filetree: &Filetree) -> String {
154        self.display_parts(filetree).join("")
155    }
156
157    #[cfg(feature = "color")]
158    pub(crate) fn treeline_color(&self, filetree: &Filetree) -> String {
159        self.display_parts_color(filetree)
160            .iter()
161            .map(ToString::to_string)
162            .collect()
163    }
164}
165
166pub(crate) struct Nodes(pub Vec<Node>);
167
168impl From<Vec<Node>> for Nodes {
169    fn from(value: Vec<Node>) -> Self {
170        Nodes(value)
171    }
172}
173
174impl From<String> for Nodes {
175    fn from(value: String) -> Self {
176        vec![Node::new(value.clone(), vec![], false)].into()
177    }
178}
179
180impl TryFrom<Vec<Value>> for Nodes {
181    type Error = Error;
182
183    fn try_from(value: Vec<Value>) -> Result<Self, Self::Error> {
184        Ok(value
185            .into_iter()
186            .map(|v| v.try_into())
187            .collect::<Result<Vec<Nodes>, Error>>()?
188            .into_iter()
189            .flat_map(|nodes| nodes.0)
190            .collect::<Vec<_>>()
191            .into())
192    }
193}
194
195impl TryFrom<serde_yaml::Mapping> for Nodes {
196    type Error = Error;
197
198    fn try_from(value: serde_yaml::Mapping) -> Result<Self, Self::Error> {
199        let mut nodes = vec![];
200        for (name, children) in value {
201            let Value::String(name) = name else {
202                return Err(Error::UnsupportedType(name.clone()));
203            };
204            let mut node = Node::new(name.clone(), vec![], true);
205            let mut child_nodes: Nodes = children.try_into()?;
206            node.add_children(&mut child_nodes.0);
207            nodes.push(node.clone());
208        }
209        Ok(nodes.into())
210    }
211}
212
213impl TryFrom<Value> for Nodes {
214    type Error = Error;
215
216    fn try_from(value: Value) -> Result<Self, Self::Error> {
217        match value {
218            Value::Null => Ok(vec![].into()),
219            Value::String(s) => Ok(s.into()),
220            Value::Sequence(seq) => seq.try_into(),
221            Value::Mapping(map) => map.try_into(),
222            v => Err(Error::UnsupportedType(v.clone())),
223        }
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn display() {
233        let node = Node::new("test".into(), vec![], true);
234        let filetree = Filetree::new(node);
235        println!("{}", filetree.root().treeline(&filetree));
236        assert_eq!(filetree.root().treeline(&filetree), "test")
237    }
238
239    #[cfg(feature = "color")]
240    #[test]
241    fn color_display() {
242        let node = Node::new("test".into(), vec![], true);
243        let filetree = Filetree::new(node);
244        println!("{}", filetree.root().treeline_color(&filetree));
245        assert_eq!(
246            filetree.root().treeline_color(&filetree),
247            "test".cyan().bold().to_string()
248        )
249    }
250}