Skip to main content

cargo_plot/core/path_view/
tree.rs

1use colored::Colorize;
2use std::collections::BTreeMap;
3use std::path::{Path, PathBuf};
4
5// Importy z rodzeństwa i innych modułów core
6use 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        // [ENG]: Logic for creating the final root node with proper weight calculation.
121        // [POL]: Logika tworzenia końcowego węzła głównego z poprawnym obliczeniem wagi.
122        let final_roots = if let Some(r_name) = root_name {
123            // [ENG]: Calculate total weight for the root node.
124            // [POL]: Obliczenie całkowitej wagi dla węzła głównego.
125            let root_bytes = if weight_cfg.dir_sum_included {
126                // [POL]: Suma wag bezpośrednich dzieci (dopasowanych elementów).
127                top_nodes.iter().map(|n| n.weight_bytes).sum()
128            } else {
129                // [POL]: Fizyczna waga folderu wejściowego z dysku.
130                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}