Added contents_first option (#19)

Added contents_first option

Added ability to yield the contents of the directory before the
directory itself

Fixes #18 

3 files changed, 110 insertions, 3 deletions

diff --git a/examples/walkdir.rs b/examples/walkdir.rs
index 4497487..c64f4ad 100644
--- a/examples/walkdir.rs
+++ b/examples/walkdir.rs
@@ -20,6 +20,7 @@ Options:
     --tree               Show output as a tree.
     --sort               Sort the output.
     -q, --ignore-errors  Ignore errors.
+    -d, --depth          Show directory's contents before the directory itself.
 ";
 
 #[derive(Debug, RustcDecodable)]
@@ -33,6 +34,7 @@ struct Args {
     flag_tree: bool,
     flag_ignore_errors: bool,
     flag_sort: bool,
+    flag_depth: bool,
 }
 
 macro_rules! wout { ($($tt:tt)*) => { {writeln!($($tt)*)}.unwrap() } }
@@ -51,6 +53,9 @@ fn main() {
     if args.flag_sort {
         walkdir = walkdir.sort_by(|a,b| a.cmp(b));
     }
+    if args.flag_depth {
+        walkdir = walkdir.contents_first(true)
+    }
     let it = walkdir.into_iter();
     let mut out = io::BufWriter::new(io::stdout());
     let mut eout = io::stderr();
diff --git a/src/lib.rs b/src/lib.rs
index ce8101e..8be0806 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -192,6 +192,7 @@ struct WalkDirOptions {
     min_depth: usize,
     max_depth: usize,
     sorter: Option<Box<FnMut(&OsString,&OsString) -> Ordering + 'static>>,
+    contents_first: bool,
 }
 
 impl WalkDir {
@@ -208,6 +209,7 @@ impl WalkDir {
                 min_depth: 0,
                 max_depth: ::std::usize::MAX,
                 sorter: None,
+                contents_first: false,
             },
             root: root.as_ref().to_path_buf(),
         }
@@ -303,6 +305,23 @@ impl WalkDir {
         self.opts.sorter = Some(Box::new(cmp));
         self
     }
+
+    /// Yield a directory's contents before the directory itself. By default, 
+    /// this is disabled.
+    ///
+    /// When `yes` is `false` (as is the default), the directory is yielded 
+    /// before its contents are read. This is useful when, e.g. you want to 
+    /// skip processing of some directories.
+    ///
+    /// When `yes` is `true`, the iterator yields the contents of a directory
+    /// before yielding the directory itself. This is useful when, e.g. you
+    /// want to recursively delete a directory.
+    pub fn contents_first(mut self, yes: bool) -> Self {
+        self.opts.contents_first = yes;
+        self
+    }
+
+
 }
 
 impl IntoIterator for WalkDir {
@@ -317,6 +336,7 @@ impl IntoIterator for WalkDir {
             stack_path: vec![],
             oldest_opened: 0,
             depth: 0,
+            deferred_dirs: vec![],
         }
     }
 }
@@ -438,6 +458,10 @@ pub struct Iter {
     /// The current depth of iteration (the length of the stack at the
     /// beginning of each iteration).
     depth: usize,
+    /// A list of DirEntries corresponding to directories, that are 
+    /// yielded after their contents has been fully yielded. This is only
+    /// used when `contents_first` is enabled.
+    deferred_dirs: Vec<DirEntry>,
 }
 
 /// A sequence of unconsumed directory entries.
@@ -506,6 +530,9 @@ impl Iterator for Iter {
         }
         while !self.stack_list.is_empty() {
             self.depth = self.stack_list.len();
+            if let Some(dentry) = self.get_deferred_dir() {
+                return Some(Ok(dentry));
+            }
             if self.depth > self.opts.max_depth {
                 // If we've exceeded the max depth, pop the current dir
                 // so that we don't descend.
@@ -523,6 +550,12 @@ impl Iterator for Iter {
                 }
             }
         }
+        if self.opts.contents_first {
+            self.depth = self.stack_list.len();
+            if let Some(dentry) = self.get_deferred_dir() {
+                return Some(Ok(dentry));
+            }
+        }
         None
     }
 }
@@ -549,7 +582,24 @@ impl Iter {
         if dent.file_type().is_dir() {
             self.push(&dent);
         }
-        if self.skippable() { None } else { Some(Ok(dent)) }
+        if dent.file_type().is_dir() && self.opts.contents_first {
+            self.deferred_dirs.push(dent);
+            None
+        } else {
+            if self.skippable() { None } else { Some(Ok(dent)) }
+        } 
+    }
+
+    fn get_deferred_dir(&mut self) -> Option<DirEntry> {
+        if self.opts.contents_first {
+            if self.depth < self.deferred_dirs.len() {
+                let deferred : DirEntry = self.deferred_dirs.pop().unwrap();
+                if !self.skippable() {
+                    return Some(deferred);
+                }
+            }
+        }
+        None
     }
 
     fn push(&mut self, dent: &DirEntry) {
diff --git a/src/tests.rs b/src/tests.rs
index e5fbf5d..6aa01b3 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1,9 +1,11 @@
 #![cfg_attr(windows, allow(dead_code, unused_imports))]
 
+use std::cmp;
 use std::env;
 use std::fs::{self, File};
 use std::io;
 use std::path::{Path, PathBuf};
+use std::collections::HashMap;
 
 use quickcheck::{Arbitrary, Gen, QuickCheck, StdGen};
 use rand::{self, Rng};
@@ -58,6 +60,41 @@ impl Tree {
         Ok(stack.pop().unwrap())
     }
 
+    fn from_walk_with_contents_first<P, F>(
+        p: P,
+        f: F,
+    ) -> io::Result<Tree>
+    where P: AsRef<Path>, F: FnOnce(WalkDir) -> WalkDir {
+        let mut contents_of_dir_at_depth = HashMap::new();
+        let mut min_depth = ::std::usize::MAX;
+        let top_level_path = p.as_ref().to_path_buf();
+        for result in f(WalkDir::new(p).contents_first(true)) {
+            let dentry = try!(result);
+
+            let tree = 
+            if dentry.file_type().is_dir() {
+                let any_contents = contents_of_dir_at_depth.remove(
+                    &(dentry.depth+1));
+            Tree::Dir(pb(dentry.file_name()), any_contents.unwrap_or_default())
+            } else {
+                if dentry.file_type().is_symlink() {
+                    let src = try!(dentry.path().read_link());
+                    let dst = pb(dentry.file_name());
+                    let dir = dentry.path().is_dir();
+                    Tree::Symlink { src: src, dst: dst, dir: dir }
+                } else {
+                    Tree::File(pb(dentry.file_name()))
+                }
+            };
+            contents_of_dir_at_depth.entry(
+                    dentry.depth).or_insert(vec!()).push(tree);
+            min_depth = cmp::min(min_depth, dentry.depth);
+        }
+        Ok(Tree::Dir(top_level_path, 
+                contents_of_dir_at_depth.remove(&min_depth)
+                .unwrap_or_default()))
+    }
+
     fn name(&self) -> &Path {
         match *self {
             Tree::Dir(ref pb, _) => pb,
@@ -299,10 +336,13 @@ fn tmpdir() -> TempDir {
 }
 
 fn dir_setup_with<F>(t: &Tree, f: F) -> (TempDir, Tree)
-        where F: FnOnce(WalkDir) -> WalkDir {
+        where F: Fn(WalkDir) -> WalkDir {
     let tmp = tmpdir();
     t.create_in(tmp.path()).unwrap();
-    let got = Tree::from_walk_with(tmp.path(), f).unwrap();
+    let got = Tree::from_walk_with(tmp.path(), &f).unwrap();
+    let got_cf = Tree::from_walk_with_contents_first(tmp.path(), &f).unwrap();
+    assert_eq!(got, got_cf);
+
     (tmp, got.unwrap_singleton().unwrap_singleton())
 }
 
@@ -556,6 +596,10 @@ fn walk_dir_min_depth_2() {
     exp.create_in(tmp.path()).unwrap();
     let got = Tree::from_walk_with(tmp.path(), |wd| wd.min_depth(2))
                    .unwrap().unwrap_dir();
+    let got_cf = Tree::from_walk_with_contents_first(
+                    tmp.path(), |wd| wd.min_depth(2))
+                   .unwrap().unwrap_dir();
+    assert_eq!(got, got_cf);
     assert_tree_eq!(exp, td("foo", got));
 }
 
@@ -571,6 +615,10 @@ fn walk_dir_min_depth_3() {
     let got = Tree::from_walk_with(tmp.path(), |wd| wd.min_depth(3))
                    .unwrap().unwrap_dir();
     assert_eq!(vec![tf("xyz")], got);
+    let got_cf = Tree::from_walk_with_contents_first(
+                    tmp.path(), |wd| wd.min_depth(3))
+                   .unwrap().unwrap_dir();
+    assert_eq!(got, got_cf);
 }
 
 #[test]
@@ -615,6 +663,10 @@ fn walk_dir_min_max_depth() {
     let got = Tree::from_walk_with(tmp.path(),
                                    |wd| wd.min_depth(2).max_depth(2))
                    .unwrap().unwrap_dir();
+    let got_cf = Tree::from_walk_with_contents_first(tmp.path(),
+                                   |wd| wd.min_depth(2).max_depth(2))
+                   .unwrap().unwrap_dir();
+    assert_eq!(got, got_cf);
     assert_tree_eq!(
         td("foo", vec![tf("bar"), td("abc", vec![]), tf("baz")]),
         td("foo", got));