extern crate ansi_term; use self::ansi_term::Colour::Fixed; use self::ansi_term::Style; use std::collections::HashSet; use utils::{ensure_end_slash, strip_end_slash_including_root}; static UNITS: [char; 4] = ['T', 'G', 'M', 'K']; pub fn draw_it( permissions: bool, short_paths: bool, depth: Option, base_dirs: HashSet, to_display: Vec<(String, u64)>, is_reversed: bool, ) { if !permissions { eprintln!("Did not have permissions for all directories"); } let mut found = HashSet::new(); let first_tree_chars = { if is_reversed { "─┴" } else { "─┬" } }; for &(ref k, _) in to_display.iter() { if base_dirs.contains(k) { display_node( &k, &mut found, &to_display, true, short_paths, depth, first_tree_chars, is_reversed, ); } } } fn get_size(nodes: &[(String, u64)], node_to_print: &str) -> Option { for &(ref k, ref v) in nodes.iter() { if *k == *node_to_print { return Some(*v); } } None } fn display_node( node: &str, nodes_already_found: &mut HashSet, to_display: &[(String, u64)], is_biggest: bool, short_paths: bool, depth: Option, indent: &str, is_reversed: bool, ) { if nodes_already_found.contains(node) { return; } nodes_already_found.insert(node.to_string()); let new_depth = match depth { None => None, Some(0) => return, Some(d) => Some(d - 1), }; match get_size(to_display, node) { None => println!("Can not find path: {}", node), Some(size) => { if !is_reversed { print_this_node(node, size, is_biggest, short_paths, indent); } fan_out( node, nodes_already_found, to_display, short_paths, new_depth, indent, is_reversed, ); if is_reversed { print_this_node(node, size, is_biggest, short_paths, indent); } } } } fn fan_out( node_to_print: &str, nodes_already_found: &mut HashSet, to_display: &[(String, u64)], short_paths: bool, new_depth: Option, indentation_str: &str, is_reversed: bool, ) { let new_indent = clean_indentation_string(indentation_str); let num_slashes = strip_end_slash_including_root(node_to_print) .matches('/') .count(); let mut num_siblings = count_siblings(to_display, num_slashes, node_to_print); let max_siblings = num_siblings; for &(ref k, _) in to_display.iter() { let temp = String::from(ensure_end_slash(node_to_print)); if k.starts_with(temp.as_str()) && k.matches('/').count() == num_slashes + 1 { num_siblings -= 1; let has_children = has_children(to_display, new_depth, k, num_slashes + 1); let (new_tree_chars, biggest) = if is_reversed { ( get_tree_chars_reversed(num_siblings != max_siblings - 1, has_children), num_siblings == 0, ) } else { ( get_tree_chars(num_siblings != 0, has_children), num_siblings == max_siblings - 1, ) }; display_node( k, nodes_already_found, to_display, biggest, short_paths, new_depth, &*(new_indent.to_string() + new_tree_chars), is_reversed, ); } } } fn clean_indentation_string(s: &str) -> String { let mut is: String = s.into(); // For reversed: is = is.replace("┌─┴", " "); is = is.replace("┌──", " "); is = is.replace("├─┴", "│ "); is = is.replace("─┴", " "); // For normal is = is.replace("└─┬", " "); is = is.replace("└──", " "); is = is.replace("├─┬", "│ "); is = is.replace("─┬", " "); // For both is = is.replace("├──", "│ "); is } fn count_siblings(to_display: &[(String, u64)], num_slashes: usize, ntp: &str) -> u64 { to_display.iter().fold(0, |a, b| { if b.0.starts_with(ntp) && b.0.as_str().matches('/').count() == num_slashes + 1 { a + 1 } else { a } }) } fn has_children( to_display: &[(String, u64)], new_depth: Option, ntp: &str, num_slashes: usize, ) -> bool { if new_depth.is_none() || new_depth.unwrap() != 1 { for &(ref k2, _) in to_display.iter() { let ntp_with_slash = String::from(ntp.to_owned() + "/"); if k2.starts_with(ntp_with_slash.as_str()) && k2.matches('/').count() == num_slashes + 1 { return true; } } } false } fn get_tree_chars(has_smaller_siblings: bool, has_children: bool) -> &'static str { if !has_smaller_siblings { if has_children { "└─┬" } else { "└──" } } else if has_children { "├─┬" } else { "├──" } } fn get_tree_chars_reversed(has_smaller_siblings: bool, has_children: bool) -> &'static str { if !has_smaller_siblings { if has_children { "┌─┴" } else { "┬──" } } else if has_children { "├─┴" } else { "├──" } } fn print_this_node( node_name: &str, size: u64, is_biggest: bool, short_paths: bool, indentation: &str, ) { let pretty_size = format!("{:>5}", human_readable_number(size),); println!( "{}", format_string( node_name, is_biggest, short_paths, pretty_size.as_ref(), indentation ) ) } pub fn format_string( dir_name: &str, is_biggest: bool, short_paths: bool, size: &str, indentation: &str, ) -> String { let printable_name = { if short_paths { dir_name.split('/').last().unwrap_or(dir_name) } else { dir_name } }; format!( "{} {} {}", if is_biggest { Fixed(196).paint(size) } else { Style::new().paint(size) }, indentation, printable_name, ) } fn human_readable_number(size: u64) -> String { for (i, u) in UNITS.iter().enumerate() { let marker = 1024u64.pow((UNITS.len() - i) as u32); if size >= marker { if size / marker < 10 { return format!("{:.1}{}", (size as f32 / marker as f32), u); } else { return format!("{}{}", (size / marker), u); } } } return format!("{}B", size); } mod tests { #[allow(unused_imports)] use super::*; #[test] fn test_human_readable_number() { assert_eq!(human_readable_number(1), "1B"); assert_eq!(human_readable_number(956), "956B"); assert_eq!(human_readable_number(1004), "1004B"); assert_eq!(human_readable_number(1024), "1.0K"); assert_eq!(human_readable_number(1536), "1.5K"); assert_eq!(human_readable_number(1024 * 512), "512K"); assert_eq!(human_readable_number(1024 * 1024), "1.0M"); assert_eq!(human_readable_number(1024 * 1024 * 1024 - 1), "1023M"); assert_eq!(human_readable_number(1024 * 1024 * 1024 * 20), "20G"); assert_eq!(human_readable_number(1024 * 1024 * 1024 * 1024), "1.0T"); } }