Merge pull request #15 from bootandy/depth2

Add support for -d depth flag

README.md

@@ -7,11 +7,11 @@ du + rust = dust. Like du but more intuitive
 
 ## Install
 
-### Cargo:
+### Cargo
 
 * cargo install du-dust
 
-### Download:
+### Download
 
 * Download linux / mac binary from [Releases](https://github.com/bootandy/dust/releases)
 * unzip file: tar -xvf _downloaded_file.tar.gz_
@@ -21,7 +21,7 @@ du + rust = dust. Like du but more intuitive
 
 Dust is meant to give you an instant overview of which directories are using disk space without requiring sort or head. Dust will print a maximum of 1 'Did not have permissions message'.
 
-Dust will list the 15 biggest sub directories or files and will smartly recurse down the tree to find the larger ones. There is no need for a '-d' flag or a '-h' flag. The largest sub directory will have its size shown in *red*
+Dust will list the 20 biggest sub directories or files and will smartly recurse down the tree to find the larger ones. There is no need for a '-d' flag or a '-h' flag. The largest sub directory will have its size shown in *red*
 
 ## Why?
 
@@ -33,10 +33,11 @@ Dust assumes that’s what you wanted to do in the first place, and takes care o
 
 ```
 Usage: dust <dir>
-Usage: dust <dir> <another_dir> <and_more>
-Usage: dust -p <dir> (full-path - does not shorten the path of the subdirectories)
-Usage: dust -s <dir> (apparent-size - shows the length of the file as opposed to the amount of disk space it uses)
-Usage: dust -n 30  <dir>  (Shows 30 directories not 15)
+Usage: dust <dir>  <another_dir> <and_more>
+Usage: dust -p <dir>  (full-path - does not shorten the path of the subdirectories)
+Usage: dust -s <dir>  (apparent-size - shows the length of the file as opposed to the amount of disk space it uses)
+Usage: dust -n 30  <dir>  (Shows 30 directories not 20)
+Usage: dust -d 3  <dir>  (Shows 3 levels of subdirectories)
 ```
 
 ```

src/display.rs

@@ -7,7 +7,8 @@ static UNITS: [char; 4] = ['T', 'G', 'M', 'K'];
 pub fn draw_it(
     permissions: bool,
     short_paths: bool,
-    base_dirs: Vec<&str>,
+    depth: Option<u64>,
+    base_dirs: Vec<String>,
     to_display: Vec<(String, u64)>,
 ) -> () {
     if !permissions {
@@ -15,13 +16,13 @@ pub fn draw_it(
     }
 
     for f in base_dirs {
-        display_node(f, &to_display, true, short_paths, "")
+        display_node(f, &to_display, true, short_paths, depth, "")
     }
 }
 
-fn get_size(nodes: &Vec<(String, u64)>, node_to_print: String) -> Option<u64> {
+fn get_size(nodes: &Vec<(String, u64)>, node_to_print: &String) -> Option<u64> {
     for &(ref k, ref v) in nodes.iter() {
-        if *k == node_to_print {
+        if *k == *node_to_print {
             return Some(*v);
         }
     }
@@ -29,18 +30,25 @@ fn get_size(nodes: &Vec<(String, u64)>, node_to_print: String) -> Option<u64> {
 }
 
 fn display_node<S: Into<String>>(
-    node_to_print: &str,
+    node_to_print: String,
     to_display: &Vec<(String, u64)>,
     is_biggest: bool,
     short_paths: bool,
+    depth: Option<u64>,
     indentation_str: S,
 ) {
-    let mut is = indentation_str.into();
-    let size = get_size(to_display, node_to_print.to_string());
-    match size {
+    let new_depth = match depth {
+        None => None,
+        Some(0) => return,
+        Some(d) => Some(d - 1),
+    };
+    match get_size(to_display, &node_to_print) {
         None => println!("Can not find path: {}", node_to_print),
         Some(size) => {
-            print_this_node(node_to_print, size, is_biggest, short_paths, is.as_ref());
+            let mut is = indentation_str.into();
+            let ntp: &str = node_to_print.as_ref();
+
+            print_this_node(ntp, size, is_biggest, short_paths, is.as_ref());
 
             is = is.replace("└─┬", "  ");
             is = is.replace("└──", "  ");
@@ -50,9 +58,7 @@ fn display_node<S: Into<String>>(
             let printable_node_slashes = node_to_print.matches('/').count();
 
             let mut num_siblings = to_display.iter().fold(0, |a, b| {
-                if b.0.starts_with(node_to_print)
-                    && b.0.matches('/').count() == printable_node_slashes + 1
-                {
+                if b.0.starts_with(ntp) && b.0.matches('/').count() == printable_node_slashes + 1 {
                     a + 1
                 } else {
                     a
@@ -61,26 +67,27 @@ fn display_node<S: Into<String>>(
 
             let mut is_biggest = true;
             for &(ref k, _) in to_display.iter() {
-                if k.starts_with(node_to_print)
-                    && k.matches('/').count() == printable_node_slashes + 1
-                {
+                if k.starts_with(ntp) && k.matches('/').count() == printable_node_slashes + 1 {
                     num_siblings -= 1;
 
                     let mut has_children = false;
-                    for &(ref k2, _) in to_display.iter() {
-                        let kk :&str = k.as_ref();
-                        if k2.starts_with(kk)
-                            && k2.matches('/').count() == printable_node_slashes + 2
-                        {
-                            has_children = true;
+                    if new_depth.is_none() || new_depth.unwrap() != 1 {
+                        for &(ref k2, _) in to_display.iter() {
+                            let kk: &str = k.as_ref();
+                            if k2.starts_with(kk)
+                                && k2.matches('/').count() == printable_node_slashes + 2
+                            {
+                                has_children = true;
+                            }
                         }
                     }
 
                     display_node(
-                        &k,
+                        k.to_string(),
                         to_display,
                         is_biggest,
                         short_paths,
+                        new_depth,
                         is.to_string() + get_tree_chars(num_siblings, has_children),
                     );
                     is_biggest = false;

src/main.rs

@@ -5,19 +5,28 @@ extern crate walkdir;
 
 use self::display::draw_it;
 use clap::{App, AppSettings, Arg};
-use utils::{find_big_ones, get_dir_tree};
+use std::io::{self, Write};
+use utils::{find_big_ones, get_dir_tree, sort};
 
 mod display;
 mod utils;
 
-static DEFAULT_NUMBER_OF_LINES: &'static str = "15";
+static DEFAULT_NUMBER_OF_LINES: &'static str = "20";
 
 fn main() {
     let options = App::new("Dust")
         .setting(AppSettings::TrailingVarArg)
+        .arg(
+            Arg::with_name("depth")
+                .short("d")
+                .long("depth")
+                .help("Depth to show")
+                .takes_value(true),
+        )
         .arg(
             Arg::with_name("number_of_lines")
                 .short("n")
+                .long("number-of-lines")
                 .help("Number of lines of output to show")
                 .takes_value(true)
                 .default_value(DEFAULT_NUMBER_OF_LINES),
@@ -25,11 +34,13 @@ fn main() {
         .arg(
             Arg::with_name("display_full_paths")
                 .short("p")
+                .long("full-paths")
                 .help("If set sub directories will not have their path shortened"),
         )
         .arg(
             Arg::with_name("display_apparent_size")
                 .short("s")
+                .long("apparent-size")
                 .help("If set will use file length. Otherwise we use blocks"),
         )
         .arg(Arg::with_name("inputs").multiple(true))
@@ -41,13 +52,46 @@ fn main() {
             Some(r) => r.collect(),
         }
     };
+
     let number_of_lines = value_t!(options.value_of("number_of_lines"), usize).unwrap();
+    let depth = {
+        if options.is_present("depth") {
+            match value_t!(options.value_of("depth"), u64) {
+                Ok(v) => Some(v + 1),
+                Err(_) => None,
+            }
+        } else {
+            None
+        }
+    };
+    if options.is_present("depth")
+        && options.value_of("number_of_lines").unwrap() != DEFAULT_NUMBER_OF_LINES
+    {
+        io::stderr()
+            .write(b"Use either -n for number of directories to show. Or -d for depth. Not both")
+            .expect("Error writing to stderr. Oh the irony!");
+        return;
+    }
+
     let use_apparent_size = options.is_present("display_apparent_size");
     let use_full_path = options.is_present("display_full_paths");
 
-    let (permissions, nodes) = get_dir_tree(&filenames, use_apparent_size);
-    let biggest_ones = find_big_ones(nodes, number_of_lines);
-    draw_it(permissions, !use_full_path, filenames, biggest_ones);
+    let (permissions, nodes, top_level_names) = get_dir_tree(&filenames, use_apparent_size);
+    let sorted_data = sort(nodes);
+    let biggest_ones = {
+        if depth.is_none() {
+            find_big_ones(sorted_data, number_of_lines)
+        } else {
+            sorted_data
+        }
+    };
+    draw_it(
+        permissions,
+        !use_full_path,
+        depth,
+        top_level_names,
+        biggest_ones,
+    );
 }
 
 #[cfg(test)]

src/tests.rs

@@ -123,60 +123,39 @@ pub fn test_soft_sym_link() {
         .output();
     assert!(c.is_ok());
 
-    let (r, r2) = soft_sym_link_output(dir_s, file_path_s, link_name_s);
+    let r = soft_sym_link_output(dir_s, file_path_s, link_name_s);
 
     // We cannot guarantee which version will appear first.
     // TODO: Consider adding predictable itteration order (sort file entries by name?)
-    let result = panic::catch_unwind(|| {
-        assert_cli::Assert::main_binary()
-            .with_args(&[dir_s])
-            .stdout()
-            .contains(r)
-            .unwrap();
-    });
-    if result.is_err() {
-        assert_cli::Assert::main_binary()
-            .with_args(&[dir_s])
-            .stdout()
-            .contains(r2)
-            .unwrap();
-    }
+    assert_cli::Assert::main_binary()
+        .with_args(&[dir_s])
+        .stdout()
+        .contains(r)
+        .unwrap();
 }
 
 #[cfg(target_os = "macos")]
-fn soft_sym_link_output(dir: &str, file_path: &str, link_name: &str) -> (String, String) {
-    let r = format!(
+fn soft_sym_link_output(dir: &str, file_path: &str, link_name: &str) -> String {
+    format!(
         "{}
 {}
 {}",
         format_string(dir, true, true, " 8.0K", ""),
         format_string(file_path, true, true, " 4.0K", "├──",),
         format_string(link_name, false, true, " 4.0K", "└──",),
-    );
-    let r2 = format!(
-        "{}
-{}
-{}",
-        format_string(dir, true, true, " 8.0K", ""),
-        format_string(link_name, true, true, " 4.0K", "├──",),
-        format_string(file_path, false, true, " 4.0K", "└──",),
-    );
-    (r, r2)
+    )
 }
 
 #[cfg(target_os = "linux")]
-fn soft_sym_link_output(dir: &str, file_path: &str, link_name: &str) -> (String, String) {
-    let r = format!(
+fn soft_sym_link_output(dir: &str, file_path: &str, link_name: &str) -> String {
+    format!(
         "{}
 {}
 {}",
         format_string(dir, true, true, " 8.0K", ""),
         format_string(file_path, true, true, " 4.0K", "├──",),
         format_string(link_name, false, true, "   0B", "└──",),
-    );
-    // I have not yet seen the output appear the other way round on linux. If this happens
-    // then add in an alterate ordering like in the mac version of this function.
-    (r, "".to_string())
+    )
 }
 
 // Hard links are ignored as the inode is the same as the file

src/utils/mod.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 use std::collections::HashSet;
+use std::cmp::Ordering;
 
 use walkdir::WalkDir;
 
@@ -9,20 +10,35 @@ use std::path::PathBuf;
 mod platform;
 use self::platform::*;
 
-pub fn get_dir_tree(filenames: &Vec<&str>, apparent_size: bool) -> (bool, HashMap<String, u64>) {
+pub fn get_dir_tree(
+    filenames: &Vec<&str>,
+    apparent_size: bool,
+) -> (bool, HashMap<String, u64>, Vec<String>) {
     let mut permissions = 0;
     let mut inodes: HashSet<(u64, u64)> = HashSet::new();
     let mut data: HashMap<String, u64> = HashMap::new();
+    let mut top_level_names = Vec::new();
+
     for b in filenames {
+        let top_level_name = strip_end_slashes(b);
         examine_dir(
-            &Path::new(b).to_path_buf(),
+            &Path::new(&top_level_name).to_path_buf(),
             apparent_size,
             &mut inodes,
             &mut data,
             &mut permissions,
         );
+        top_level_names.push(top_level_name);
     }
-    (permissions == 0, data)
+    (permissions == 0, data, top_level_names)
+}
+
+fn strip_end_slashes(s: &str) -> String {
+    let mut new_name = String::from(s);
+    while new_name.chars().last() == Some('/') && new_name.len() != 1 {
+        new_name.pop();
+    }
+    new_name
 }
 
 fn examine_dir(
@@ -65,12 +81,23 @@ fn examine_dir(
         }
     }
 }
+pub fn compare_tuple(a :&(String, u64), b: &(String, u64)) -> Ordering {
+    let result = b.1.cmp(&a.1);
+    if result == Ordering::Equal {
+        a.0.cmp(&b.0)
+    } else {
+        result
+    }
+}
 
-pub fn find_big_ones<'a>(data: HashMap<String, u64>, max_to_show: usize) -> Vec<(String, u64)> {
+pub fn sort<'a>(data: HashMap<String, u64>) -> Vec<(String, u64)> {
     let mut new_l: Vec<(String, u64)> = data.iter().map(|(a, b)| (a.clone(), *b)).collect();
-    new_l.sort_by(|a, b| b.1.cmp(&a.1));
+    new_l.sort_by(|a, b| compare_tuple(&a, &b));
+    new_l
+}
 
-    if new_l.len() > max_to_show {
+pub fn find_big_ones<'a>(new_l: Vec<(String, u64)>, max_to_show: usize) -> Vec<(String, u64)> {
+    if max_to_show > 0 && new_l.len() > max_to_show {
         new_l[0..max_to_show + 1].to_vec()
     } else {
         new_l