Skip to main content

ferrijs_std/utils/
fs.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3use std::{fs::Metadata, io, path::PathBuf};
4
5use tokio::fs::{self};
6
7pub struct DirectoryWalker<T>
8where
9    T: Fn(&str) -> bool,
10{
11    stack: Vec<(PathBuf, Option<Metadata>)>,
12    filter: T,
13    recursive: bool,
14    eat_root: bool,
15}
16
17impl<T> DirectoryWalker<T>
18where
19    T: Fn(&str) -> bool,
20{
21    pub fn new(root: PathBuf, filter: T) -> Self {
22        Self {
23            stack: vec![(root, None)],
24            filter,
25            recursive: false,
26            eat_root: true,
27        }
28    }
29
30    pub fn set_recursive(&mut self, recursive: bool) {
31        self.recursive = recursive;
32    }
33
34    pub async fn walk(&mut self) -> io::Result<Option<(PathBuf, Metadata)>> {
35        if self.eat_root {
36            self.eat_root = false;
37            let (dir, _) = self.stack.pop().unwrap();
38            self.append_stack(&dir).await?;
39        }
40        if let Some((entry, metadata)) = self.stack.pop() {
41            let metadata = metadata.unwrap();
42            if self.recursive && metadata.is_dir() {
43                self.append_stack(&entry).await?;
44            }
45
46            Ok(Some((entry, metadata)))
47        } else {
48            Ok(None)
49        }
50    }
51
52    pub fn walk_sync(&mut self) -> io::Result<Option<(PathBuf, Metadata)>> {
53        if self.eat_root {
54            self.eat_root = false;
55            let (dir, _) = self.stack.pop().unwrap();
56            self.append_stack_sync(&dir)?;
57        }
58        if let Some((entry, metadata)) = self.stack.pop() {
59            let metadata = metadata.unwrap();
60            if self.recursive && metadata.is_dir() {
61                self.append_stack_sync(&entry)?;
62            }
63
64            Ok(Some((entry, metadata)))
65        } else {
66            Ok(None)
67        }
68    }
69
70    async fn append_stack(&mut self, dir: &PathBuf) -> io::Result<()> {
71        let mut stream = fs::read_dir(dir).await?;
72
73        while let Some(entry) = stream.next_entry().await? {
74            let name = entry.file_name();
75            let name = name.to_string_lossy();
76            if !(self.filter)(name.as_ref()) {
77                continue;
78            }
79            let entry_path = entry.path();
80            let metadata = fs::symlink_metadata(&entry_path).await?;
81
82            self.stack.push((entry_path, Some(metadata)));
83        }
84        Ok(())
85    }
86
87    fn append_stack_sync(&mut self, dir: &PathBuf) -> io::Result<()> {
88        let dir = std::fs::read_dir(dir)?;
89
90        for entry in dir.flatten() {
91            let name = entry.file_name();
92            let name = name.to_string_lossy();
93            if !(self.filter)(name.as_ref()) {
94                continue;
95            }
96            let entry_path = entry.path();
97            let metadata = entry_path.symlink_metadata()?;
98            self.stack.push((entry_path, Some(metadata)))
99        }
100
101        Ok(())
102    }
103}