diff --git a/Cargo.toml b/Cargo.toml index 43ea76d..1aa13c1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "async-walkdir" -version = "2.2.0" +version = "2.2.1" authors = ["Ririsoft "] edition = "2021" description = "Asynchronous directory traversal for Rust." diff --git a/src/lib.rs b/src/lib.rs index 97b4a0f..38c541a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -84,6 +84,7 @@ mod error; use std::future::Future; +use std::marker::PhantomData; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::task::{Context, Poll}; @@ -112,9 +113,22 @@ type BoxStream = futures_lite::stream::Boxed>; /// # Panics /// /// Panics if the directories depth overflows `usize`. -pub struct WalkDir { +pub struct WalkDir { root: PathBuf, entries: BoxStream, + opts: WalkDirOptions, + _state: PhantomData, +} + +/// A marker struct to signal, that the DirWalk is not being filtered +pub struct Unfiltered; +/// A marker struct to snignal, that the DirWalk has a filter applied. +/// This results in some options not being configurable after this point anymore. +pub struct Filtered; + +#[derive(Default, Clone, Copy)] +struct WalkDirOptions { + contents_first: bool, } /// Sets the filtering behavior. @@ -136,26 +150,58 @@ impl WalkDir { root: root.as_ref().to_owned(), entries: walk_dir( root, + WalkDirOptions::default(), + None:: BoxedFut + Send>>, + ), + opts: WalkDirOptions::default(), + _state: PhantomData, + } + } +} +impl WalkDir { + /// Yield the directory's content before the directory. + /// + /// When `yes` is false (the default), the directory is yielded before the contents are read. + /// This is useful when e.g. you want to skip processing some of the 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 { + let root = self.root.clone(); + self.opts.contents_first = yes; + Self { + root: self.root, + entries: walk_dir( + root, + self.opts, None:: BoxedFut + Send>>, ), + opts: self.opts, + _state: PhantomData, } } /// Filter entries. - pub fn filter(self, f: F) -> Self + pub fn filter(self, f: F) -> WalkDir where F: FnMut(DirEntry) -> Fut + Send + 'static, Fut: Future + Send, { let root = self.root.clone(); - Self { + WalkDir { root: self.root, - entries: walk_dir(root, Some(f)), + entries: walk_dir(root, self.opts, Some(f)), + opts: self.opts, + _state: PhantomData, } } } -impl Stream for WalkDir { +impl Stream for WalkDir +where + State: Unpin, +{ type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { @@ -164,7 +210,20 @@ impl Stream for WalkDir { } } -fn walk_dir(root: impl AsRef, filter: Option) -> BoxStream +struct DirStackItem { + path: PathBuf, + read_dir: ReadDir, + // holds the directory, while it is being traversed + // only relevant when `contents_first` is set to true + deferred_entry: Option, + // when `contents_first` is enabled, the file tree + // is traversed in post-order, so all non-directory + // contents of a directory need to be stored in here + // until the child directories are finished traversing + deferred_files: Vec, +} + +fn walk_dir(root: impl AsRef, opts: WalkDirOptions, filter: Option) -> BoxStream where F: FnMut(DirEntry) -> Fut + Send + 'static, Fut: Future + Send, @@ -178,9 +237,21 @@ where Err(InnerError::Io { path: root, source }.into()), State::Done, )), - Ok(rd) => walk(vec![(root, rd)], filter).await, + Ok(rd) => { + walk( + vec![DirStackItem { + path: root, + read_dir: rd, + deferred_files: Vec::new(), + deferred_entry: None, + }], + filter, + opts, + ) + .await + } }, - State::Walk((dirs, filter)) => walk(dirs, filter).await, + State::Walk((dirs, filter)) => walk(dirs, filter, opts).await, State::Done => None, } }, @@ -190,7 +261,7 @@ where enum State { Start((PathBuf, Option)), - Walk((Vec<(PathBuf, ReadDir)>, Option)), + Walk((Vec, Option)), Done, } @@ -198,41 +269,106 @@ type UnfoldState = (Result, State); // Iterative on purpose: filtered-out entries must not grow the stack (see issue #13). async fn walk( - mut dirs: Vec<(PathBuf, ReadDir)>, + mut dirs: Vec, mut filter: Option, + opts: WalkDirOptions, ) -> Option> where F: FnMut(DirEntry) -> Fut + Send + 'static, Fut: Future + Send, { loop { - let (path, dir) = dirs.last_mut()?; - let entry = match dir.next().await { + let res = { + let item = dirs.last_mut()?; + item.read_dir.next().await + }; + + let entry = match res { Some(Ok(entry)) => entry, - Some(Err(source)) => return io_error(path.to_path_buf(), source, dirs, filter), + Some(Err(source)) => { + let path = dirs.last().unwrap().path.clone(); + return io_error(path, source, dirs, filter); + } None => { - dirs.pop(); - continue; + let Some(mut popped) = dirs.pop() else { + unreachable!() + }; + if let Some(deferred_file) = popped.deferred_files.pop() { + dirs.push(popped); + return Some((Ok(deferred_file), State::Walk((dirs, filter)))); + } else if let Some(deferred_dir) = popped.deferred_entry.take() { + return Some((Ok(deferred_dir), State::Walk((dirs, filter)))); + } else { + continue; + } } }; + let ft = match entry.file_type().await { Ok(ft) => ft, Err(source) => return io_error(entry.path(), source, dirs, filter), }; + let filtering = match filter.as_mut() { Some(filter) => filter(entry.clone()).await, None => Filtering::Continue, }; - if ft.is_dir() && filtering != Filtering::IgnoreDir { - let path = entry.path(); - let rd = match read_dir(&path).await { - Ok(rd) => rd, - Err(source) => return io_error(path, source, dirs, filter), - }; - dirs.push((path, rd)); - } - if filtering == Filtering::Continue { - return Some((Ok(entry), State::Walk((dirs, filter)))); + + match filtering { + Filtering::IgnoreDir => continue, + Filtering::Ignore => { + if ft.is_dir() { + let path = entry.path(); + let rd = match read_dir(&path).await { + Ok(rd) => rd, + Err(source) => return io_error(path, source, dirs, filter), + }; + dirs.push(DirStackItem { + path, + read_dir: rd, + deferred_entry: None, + deferred_files: Vec::new(), + }); + } + continue; + } + Filtering::Continue => { + if ft.is_dir() { + let path = entry.path(); + let rd = match read_dir(&path).await { + Ok(rd) => rd, + Err(source) => return io_error(path, source, dirs, filter), + }; + + let deferred_entry = if opts.contents_first { + Some(entry.clone()) + } else { + None + }; + + dirs.push(DirStackItem { + path, + read_dir: rd, + deferred_entry, + deferred_files: Vec::new(), + }); + + if opts.contents_first { + continue; + } else { + return Some((Ok(entry), State::Walk((dirs, filter)))); + } + } else { + if opts.contents_first { + if let Some(last) = dirs.last_mut() { + last.deferred_files.push(entry); + } + continue; + } else { + return Some((Ok(entry), State::Walk((dirs, filter)))); + } + } + } } } } @@ -240,7 +376,7 @@ where fn io_error( path: PathBuf, source: std::io::Error, - dirs: Vec<(PathBuf, ReadDir)>, + dirs: Vec, filter: Option, ) -> Option> { let err = InnerError::Io { path, source }.into(); @@ -390,6 +526,40 @@ mod tests { Ok(()) }) } + #[test] + fn contents_first() -> Result<()> { + block_on(async { + let root = tempfile::tempdir()?; + let f1 = root.path().join("f1.txt"); + let d1 = root.path().join("d1"); + let f2 = d1.join("f2.txt"); + let d2 = d1.join("d2"); + let f3 = d2.join("f3.txt"); + + async_fs::create_dir_all(&d2).await?; + async_fs::write(&f1, []).await?; + async_fs::write(&f2, []).await?; + async_fs::write(&f3, []).await?; + + let want = vec![ + f3.to_owned(), + d2.to_owned(), + f2.to_owned(), + d1.to_owned(), + f1.to_owned(), + ]; + let mut wd = WalkDir::new(root.path()).contents_first(true); + + let mut got = Vec::new(); + while let Some(entry) = wd.next().await { + let entry = entry.unwrap(); + got.push(entry.path()); + } + assert_eq!(got, want); + + Ok(()) + }) + } #[test] fn filter_dirs_no_traverse() -> Result<()> {