Merge pull request #12 from alirezasariri78/dev

Dev

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "DiskAnalyzer"
-version = "0.2.2"
+version = "0.2.3"
 edition = "2021"
 repository="https://github.com/alirezasariri78/DiskAnalyzer"
 description="make diagram base on system folder size"
@@ -17,3 +17,4 @@ path = "src/main.rs"
 clap={version="4.5.4",features=["cargo"]}
 rs-abbreviation-number = "0.2.1"
 cli-table = "0.4"
+colored = "2.1.0"

src/diagram.rs

@@ -1,9 +1,13 @@
-use cli_table::print_stdout;
+mod table;
+mod tree;
+mod shared;
 
+
+use cli_table::print_stdout;
 use crate::crawler::Node;
 use std::sync::Arc;
-mod table;
-mod tree;
+
+
 use crate::args::{CommandArgs, DiagramType};
 pub fn show_diagram(root: &Arc<Node>, arguments: &CommandArgs) {
     match arguments.diagram {

src/diagram/shared.rs

@@ -0,0 +1,11 @@
+
+pub(crate) fn get_color_from_size(size:u64)->&'static str{
+    const GREEN_SIZE_BYTES:u64=1_024_000_000;
+    const YELLOW_SIZE_BYTES:u64=GREEN_SIZE_BYTES*10;
+    const YELLOW_START:u64=GREEN_SIZE_BYTES+1; 
+    match size {
+        0..=GREEN_SIZE_BYTES=>"Green",
+        YELLOW_START..=YELLOW_SIZE_BYTES=>"Yellow",
+        _=>"Red",
+    }
+}

src/diagram/table.rs

@@ -3,10 +3,14 @@ use crate::{
     args::{CommandArgs, SortType},
     util::thousends_seperator,
 };
+use cli_table::Color;
 use cli_table::{format::Justify, Cell, CellStruct, Style, Table, TableStruct};
 use rs_abbreviation_number::NumericAbbreviate;
+use std::str::FromStr;
 use std::{ops::Deref, sync::Arc};
 
+use super::shared::get_color_from_size;
+
 pub fn create_table_diagram(tree: &Arc<Node>, args: &CommandArgs) -> TableStruct {
     let mut nodes: Vec<Arc<Node>> = Vec::new();
     crawl_tree(tree, args, &mut nodes);
@@ -23,14 +27,15 @@ fn export_to_table(nodes: Vec<Arc<Node>>, args: &CommandArgs) -> Vec<Vec<CellStr
     let mut result: Vec<Vec<CellStruct>> = Vec::new();
     let sorted_data = sort_data(&nodes, args.sort.clone());
     for (name, size) in sorted_data {
-        let size = format!(
+        let formated_size = format!(
             "{}iB ({} bytes)",
             size.abbreviate_number(),
             thousends_seperator(size)
         );
+        let color=Color::from_str(get_color_from_size(size)).unwrap_or(Color::White);
         result.push(vec![
-            name.cell().justify(Justify::Center),
-            size.cell().justify(Justify::Center),
+            name.cell().foreground_color(Some(color)).justify(Justify::Center),
+            formated_size.cell().foreground_color(Some(color)).justify(Justify::Center),
         ]);
     }
     result

src/diagram/tree.rs

@@ -1,9 +1,10 @@
-use rs_abbreviation_number::NumericAbbreviate;
-
 use crate::args::CommandArgs;
 use crate::crawler::Node;
+use crate::diagram::shared::*;
 use crate::util::*;
-use std::{ops::Deref, sync::Arc};
+use colored::{Color, Colorize};
+use rs_abbreviation_number::NumericAbbreviate;
+use std::{ops::Deref, str::FromStr, sync::Arc};
 
 pub const MIDDLE_CHAR: &'static str = "├──";
 pub const END_CHAR: &'static str = "└──";
@@ -31,11 +32,7 @@ fn crawl_tree(tree: &Arc<Node>, args: &CommandArgs, result: &mut String) {
 }
 
 fn add_branch(node: &Node) -> String {
-    let mut branch_char = MIDDLE_CHAR;
-    if node.is_last_child() {
-        branch_char = END_CHAR;
-    }
-
+    let branch_char = get_branch_char(node);
     let mut size = node.get_size().get().abbreviate_number();
     size.push_str("iB");
     format!(
@@ -51,3 +48,52 @@ fn add_branch(node: &Node) -> String {
         )
     )
 }
+
+fn get_branch_char(node: &Node) -> String {
+    let mut branch_char = MIDDLE_CHAR;
+    if node.is_last_child() {
+        branch_char = END_CHAR;
+    }
+    let color_str = get_color_from_size(node.get_size().get());
+    let colored = Color::from_str(color_str).unwrap_or(Color::White);
+    branch_char.color(colored).to_string()
+}
+
+mod tests {
+
+    use std::path::PathBuf;
+
+    use super::*;
+    #[test]
+    fn get_branch_char_heavy_size_test() {
+        let node = Node::new(PathBuf::from(""), "root".to_string(), 1_000_000_000_000, 0);
+        assert_eq!("\u{1b}[31m├──\u{1b}[0m", get_branch_char(&node));
+    }
+    #[test]
+    fn get_branch_char_medium_size_test() {
+        let node = Node::new(PathBuf::from(""), "root".to_string(), 8_024_000_000, 0);
+        assert_eq!("\u{1b}[33m├──\u{1b}[0m", get_branch_char(&node));
+    }
+
+    #[test]
+    fn get_branch_char_small_size_test() {
+        let node = Node::new(PathBuf::from(""), "root".to_string(), 5, 0);
+        assert_eq!("\u{1b}[32m├──\u{1b}[0m", get_branch_char(&node));
+    }
+
+    #[test]
+    fn add_branch_root_test() {
+        let node = Node::new(PathBuf::from(""), "root".to_string(), 5000, 0);
+        let input = add_branch(&node);
+        let expect = "\n\u{1b}[32m├──\u{1b}[0mroot 5KiB(5,000 bytes)".to_string();
+        assert_eq!(expect, input);
+    }
+
+    #[test]
+    fn add_branch_child_test() {
+        let child = Node::new(PathBuf::from(""), "child".to_string(), 8_024_000_000, 1);
+        let input = add_branch(&child);
+        let expect = "\n\t\u{1b}[33m├──\u{1b}[0mchild 8.02GiB(8,024,000,000 bytes)".to_string();
+        assert_eq!(expect, input);
+    }
+}

src/util/mod.rs

@@ -1,6 +1,6 @@
 
 
-pub fn thousends_seperator(i: u64) -> String {
+pub(crate) fn thousends_seperator(i: u64) -> String {
     let mut s = String::new();
     let i_str = i.to_string();
     let a = i_str.chars().rev().enumerate();