Add support for -d depth flag

Following a user request the option '-d N' allows a user to output N
levels of sub directories.

Fixed bug: so that trailing slashes are now removed.

src/main.rs

@@ -5,7 +5,8 @@ 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;
@@ -15,6 +16,12 @@ static DEFAULT_NUMBER_OF_LINES: &'static str = "15";
 fn main() {
     let options = App::new("Dust")
         .setting(AppSettings::TrailingVarArg)
+        .arg(
+            Arg::with_name("depth")
+                .short("d")
+                .help("Depth to show")
+                .takes_value(true),
+        )
         .arg(
             Arg::with_name("number_of_lines")
                 .short("n")
@@ -41,13 +48,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)]