Added a filter to the directory iterator

diritertest/src/diriter.rs

@@ -1,21 +1,39 @@
 //! Iterator over directories
 
 use std::collections::VecDeque;
-use std::fs::DirEntry;
 use std::fs::ReadDir;
+use std::marker::PhantomData;
 use std::path::Path;
 use std::path::PathBuf;
 
-pub struct DirIter {
+pub trait Accept {
+    /// Accept a directory and return it as a result. Note that this stops the processing of any
+    /// other entries in the directory.
+    fn accept_dir(_: &Path) -> bool {
+        false
+    }
+
+    /// Accept a file and return it as a result.
+    fn accept_file(_: &Path) -> bool {
+        false
+    }
+}
+
+pub struct DirIter<A: Accept> {
     curr_dir: ReadDir,
     remaining_dirs: VecDeque<PathBuf>,
+    accept: PhantomData<A>,
 }
 
-impl DirIter {
+impl<A> DirIter<A>
+where
+    A: Accept,
+{
     pub fn new(start: &Path) -> std::io::Result<Self> {
         Ok(Self {
             curr_dir: start.read_dir()?,
             remaining_dirs: VecDeque::new(),
+            accept: PhantomData,
         })
     }
 
@@ -23,23 +41,32 @@ impl DirIter {
         let next = self.curr_dir.next();
         match next {
             Some(Ok(x)) => {
-                if x.path().is_dir() {
-                    self.remaining_dirs.push_back(x.path());
+                let x = x.path();
+                if x.is_dir() {
+                    if A::accept_dir(&x) {
+                        Some(x)
+                    } else {
+                        self.remaining_dirs.push_back(x);
                         self.advance()
+                    }
+                } else {
+                    if A::accept_file(&x) {
+                        Some(x)
                     } else {
-                    Some(x.path())
+                        self.advance()
+                    }
                 }
             }
-            Some(Err(_)) | None => {
+            _ => {
                 let next_dir = self.remaining_dirs.pop_front()?;
-                self.curr_dir = std::fs::read_dir(next_dir).ok()?;
+                self.curr_dir = next_dir.read_dir().ok()?;
                 self.advance()
             }
         }
     }
 }
 
-impl Iterator for DirIter {
+impl<A: Accept> Iterator for DirIter<A> {
     type Item = PathBuf;
 
     fn next(&mut self) -> Option<Self::Item> {

diritertest/src/main.rs

@@ -1,21 +1,39 @@
+use std::ffi::OsStr;
 use std::path::Path;
 
-use crate::diriter::DirIter;
+use diriter::Accept;
+
+use diriter::DirIter;
 
 mod diriter;
 
+struct AllFiles;
+impl Accept for AllFiles {
+    fn accept_file(_: &Path) -> bool {
+        true
+    }
+    fn accept_dir(_: &Path) -> bool {
+        false
+    }
+}
+
+struct Sources;
+impl Accept for Sources {
+    fn accept_dir(p: &Path) -> bool {
+        p.file_name().unwrap() == OsStr::new("src")
+    }
+}
+
 fn main() {
-    let mut all_iter = DirIter::new(&Path::new(".")).unwrap();
+    let mut all_iter = DirIter::<AllFiles>::new(&Path::new(".")).unwrap();
     while let Some(path) = all_iter.next() {
         println!(">> {:?}", path);
     }
 
-    println!("-----------------------");
+    println!("{}", std::iter::repeat('-').take(80).collect::<String>());
 
-    let all_iter = DirIter::new(&Path::new(".")).unwrap();
-    for path in all_iter {
-        println!(">>> {:?}", path);
+    let iter = DirIter::<Sources>::new(&Path::new(".")).unwrap();
+    for name in iter {
+        println!(">> {:?}", name);
     }
-
-    println!("Hello, world!");
 }