cargo_plot/core/path_view/
tree.rs1use colored::Colorize;
2use std::collections::BTreeMap;
3use std::path::{Path, PathBuf};
4
5use super::node::FileNode;
7use crate::core::file_stats::weight::{self, WeightConfig};
8use crate::core::path_matcher::SortStrategy;
9use crate::theme::for_path_tree::{DIR_ICON, FILE_ICON, TreeStyle, get_file_type};
10pub struct PathTree {
11 roots: Vec<FileNode>,
12 style: TreeStyle,
13}
14
15impl PathTree {
16 #[must_use]
17 pub fn build(
18 paths_strings: &[String],
19 base_dir: &str,
20 sort_strategy: SortStrategy,
21 weight_cfg: &WeightConfig,
22 root_name: Option<&str>,
23 no_emoji: bool,
24 ) -> Self {
25 let base_path_obj = Path::new(base_dir);
26 let paths: Vec<PathBuf> = paths_strings.iter().map(PathBuf::from).collect();
27 let mut tree_map: BTreeMap<PathBuf, Vec<PathBuf>> = BTreeMap::new();
28
29 for p in &paths {
30 let parent = p
31 .parent()
32 .map_or_else(|| PathBuf::from("."), Path::to_path_buf);
33 tree_map.entry(parent).or_default().push(p.clone());
34 }
35
36 fn build_node(
37 path: &PathBuf,
38 paths_map: &BTreeMap<PathBuf, Vec<PathBuf>>,
39 base_path: &Path,
40 sort_strategy: SortStrategy,
41 weight_cfg: &WeightConfig,
42 no_emoji: bool,
43 ) -> FileNode {
44 let name = path
45 .file_name()
46 .map_or_else(|| "/".to_string(), |n| n.to_string_lossy().to_string());
47
48 let is_dir = path.is_dir() || path.to_string_lossy().ends_with('/');
49 let icon = if no_emoji {
50 String::new()
51 } else if is_dir {
52 DIR_ICON.to_string()
53 } else if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
54 get_file_type(ext).icon.to_string()
55 } else {
56 FILE_ICON.to_string()
57 };
58
59 let absolute_path = base_path.join(path);
60 let mut weight_bytes =
61 weight::get_path_weight(&absolute_path, weight_cfg.dir_sum_included);
62 let mut children = vec![];
63
64 if let Some(child_paths) = paths_map.get(path) {
65 let mut child_nodes: Vec<FileNode> = child_paths
66 .iter()
67 .map(|c| {
68 build_node(c, paths_map, base_path, sort_strategy, weight_cfg, no_emoji)
69 })
70 .collect();
71
72 FileNode::sort_slice(&mut child_nodes, sort_strategy);
73
74 if is_dir && weight_cfg.dir_sum_included {
75 weight_bytes = child_nodes.iter().map(|n| n.weight_bytes).sum();
76 }
77 children = child_nodes;
78 }
79
80 let weight_str = weight::format_weight(weight_bytes, is_dir, weight_cfg);
81
82 FileNode {
83 name,
84 path: path.clone(),
85 is_dir,
86 icon,
87 weight_str,
88 weight_bytes,
89 children,
90 }
91 }
92
93 let roots_paths: Vec<PathBuf> = paths
94 .iter()
95 .filter(|p| {
96 let parent = p.parent();
97 parent.is_none()
98 || parent.unwrap() == Path::new("")
99 || !paths.contains(&parent.unwrap().to_path_buf())
100 })
101 .cloned()
102 .collect();
103
104 let mut top_nodes: Vec<FileNode> = roots_paths
105 .into_iter()
106 .map(|r| {
107 build_node(
108 &r,
109 &tree_map,
110 base_path_obj,
111 sort_strategy,
112 weight_cfg,
113 no_emoji,
114 )
115 })
116 .collect();
117
118 FileNode::sort_slice(&mut top_nodes, sort_strategy);
119
120 let final_roots = if let Some(r_name) = root_name {
123 let root_bytes = if weight_cfg.dir_sum_included {
126 top_nodes.iter().map(|n| n.weight_bytes).sum()
128 } else {
129 weight::get_path_weight(base_path_obj, false)
131 };
132
133 let root_weight_str = weight::format_weight(root_bytes, true, weight_cfg);
134
135 vec![FileNode {
136 name: r_name.to_string(),
137 path: PathBuf::from(r_name),
138 is_dir: true,
139 icon: if no_emoji {
140 String::new()
141 } else {
142 DIR_ICON.to_string()
143 },
144 weight_str: root_weight_str,
145 weight_bytes: root_bytes,
146 children: top_nodes,
147 }]
148 } else {
149 top_nodes
150 };
151
152 Self {
153 roots: final_roots,
154 style: TreeStyle::default(),
155 }
156 }
157
158 #[must_use]
159 pub fn render_cli(&self) -> String {
160 self.plot(&self.roots, "", true)
161 }
162
163 #[must_use]
164 pub fn render_txt(&self) -> String {
165 self.plot(&self.roots, "", false)
166 }
167
168 fn plot(&self, nodes: &[FileNode], indent: &str, use_color: bool) -> String {
169 let mut result = String::new();
170 for (i, node) in nodes.iter().enumerate() {
171 let is_last = i == nodes.len() - 1;
172 let has_children = !node.children.is_empty();
173
174 let branch = if node.is_dir {
175 match (is_last, has_children) {
176 (true, true) => &self.style.dir_last_with_children,
177 (false, true) => &self.style.dir_mid_with_children,
178 (true, false) => &self.style.dir_last_no_children,
179 (false, false) => &self.style.dir_mid_no_children,
180 }
181 } else if is_last {
182 &self.style.file_last
183 } else {
184 &self.style.file_mid
185 };
186
187 let weight_prefix = if node.weight_str.is_empty() {
188 String::new()
189 } else if use_color {
190 node.weight_str.truecolor(120, 120, 120).to_string()
191 } else {
192 node.weight_str.clone()
193 };
194
195 let line = if use_color {
196 if node.is_dir {
197 format!(
198 "{weight_prefix}{}{branch_color} {icon} {name}\n",
199 indent.green(),
200 branch_color = branch.green(),
201 icon = node.icon,
202 name = node.name.truecolor(200, 200, 50)
203 )
204 } else {
205 format!(
206 "{weight_prefix}{}{branch_color} {icon} {name}\n",
207 indent.green(),
208 branch_color = branch.green(),
209 icon = node.icon,
210 name = node.name.white()
211 )
212 }
213 } else {
214 format!(
215 "{weight_prefix}{indent}{branch} {icon} {name}\n",
216 icon = node.icon,
217 name = node.name
218 )
219 };
220
221 result.push_str(&line);
222
223 if has_children {
224 let new_indent = if is_last {
225 format!("{indent}{}", self.style.indent_last)
226 } else {
227 format!("{indent}{}", self.style.indent_mid)
228 };
229 result.push_str(&self.plot(&node.children, &new_indent, use_color));
230 }
231 }
232 result
233 }
234}