anti_fs/
du.rs

1use std::fs::{self, ReadDir};
2use std::io;
3use std::path::{Path, PathBuf};
4
5#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize)]
6pub enum FormatUnit {
7    Bytes,
8    Kilo,
9    Mega,
10    Giga,
11    Human,
12}
13
14#[derive(Debug, Clone)]
15pub struct ParseFormatUnitError(pub String);
16
17impl std::fmt::Display for ParseFormatUnitError {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        write!(f, "{}", self.0)
20    }
21}
22
23impl std::error::Error for ParseFormatUnitError {}
24
25impl std::str::FromStr for FormatUnit {
26    type Err = ParseFormatUnitError;
27
28    fn from_str(s: &str) -> Result<Self, Self::Err> {
29        match s.to_ascii_lowercase().as_str() {
30            "bytes" | "b" => Ok(FormatUnit::Bytes),
31            "kilo" | "kb" | "k" => Ok(FormatUnit::Kilo),
32            "mega" | "mb" | "m" => Ok(FormatUnit::Mega),
33            "giga" | "gb" | "g" => Ok(FormatUnit::Giga),
34            "human" | "h" | "auto" => Ok(FormatUnit::Human),
35            _ => Err(ParseFormatUnitError(format!("invalid format unit: {s}"))),
36        }
37    }
38}
39
40#[derive(Debug, Clone)]
41pub struct DuOptions {
42    pub depth: usize,
43    pub unit: FormatUnit,
44}
45
46impl Default for DuOptions {
47    fn default() -> Self {
48        Self {
49            depth: usize::MAX,
50            unit: FormatUnit::Human,
51        }
52    }
53}
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub struct Entry {
57    pub path: PathBuf,
58    pub size: u64,
59}
60
61pub fn du<P: AsRef<Path>>(path: P, opts: &DuOptions) -> io::Result<Vec<Entry>> {
62    let mut results = Vec::new();
63    du_inner(path.as_ref(), 0, opts, &mut results)?;
64    Ok(results)
65}
66
67fn du_inner(
68    path: &Path,
69    current: usize,
70    opts: &DuOptions,
71    results: &mut Vec<Entry>,
72) -> io::Result<u64> {
73    let metadata = fs::symlink_metadata(path)?;
74    if metadata.is_file() {
75        let size = metadata.len();
76        if current <= opts.depth {
77            results.push(Entry {
78                path: path.to_path_buf(),
79                size,
80            });
81        }
82        return Ok(size);
83    }
84
85    let mut total = 0u64;
86    if metadata.is_dir() {
87        let entries: ReadDir = fs::read_dir(path)?;
88        for entry in entries {
89            let entry = entry?;
90            total += du_inner(&entry.path(), current + 1, opts, results)?;
91        }
92        if current <= opts.depth {
93            results.push(Entry {
94                path: path.to_path_buf(),
95                size: total,
96            });
97        }
98    }
99    Ok(total)
100}
101
102pub fn format_size(bytes: u64, unit: FormatUnit) -> String {
103    match unit {
104        FormatUnit::Bytes => bytes.to_string(),
105        FormatUnit::Kilo => format!("{:.1}KB", bytes as f64 / 1024.0),
106        FormatUnit::Mega => format!("{:.1}MB", bytes as f64 / 1024.0 / 1024.0),
107        FormatUnit::Giga => format!("{:.1}GB", bytes as f64 / 1024.0 / 1024.0 / 1024.0),
108        FormatUnit::Human => {
109            const UNITS: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
110            let mut size = bytes as f64;
111            let mut unit_idx = 0usize;
112            while size >= 1024.0 && unit_idx < UNITS.len() - 1 {
113                size /= 1024.0;
114                unit_idx += 1;
115            }
116            format!("{:.1}{}", size, UNITS[unit_idx])
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use tempfile::tempdir;
125
126    #[test]
127    fn test_du_includes_files() {
128        let dir = tempdir().unwrap();
129        let file_path = dir.path().join("file.txt");
130        fs::write(&file_path, b"hello").unwrap();
131
132        let opts = DuOptions {
133            depth: 0,
134            unit: FormatUnit::Bytes,
135        };
136        let results = du(&file_path, &opts).unwrap();
137        assert_eq!(results.len(), 1);
138        assert_eq!(results[0].size, 5);
139    }
140
141    #[test]
142    fn test_du_nested_depth() {
143        let dir = tempdir().unwrap();
144        let sub = dir.path().join("sub");
145        fs::create_dir(&sub).unwrap();
146        fs::write(sub.join("a"), b"1234567890").unwrap();
147        fs::write(dir.path().join("b"), b"1234").unwrap();
148
149        let opts = DuOptions {
150            depth: 1,
151            unit: FormatUnit::Bytes,
152        };
153        let mut results = du(dir.path(), &opts).unwrap();
154        results.sort_by(|a, b| a.path.cmp(&b.path));
155        assert_eq!(results.len(), 3); // root, sub, file b
156        let root = results.iter().find(|e| e.path == dir.path()).unwrap();
157        assert_eq!(root.size, 14);
158    }
159
160    #[test]
161    fn test_format_size_units() {
162        assert_eq!(format_size(1024, FormatUnit::Human), "1.0KB");
163        assert_eq!(format_size(1024, FormatUnit::Kilo), "1.0KB");
164        assert_eq!(format_size(1024 * 1024, FormatUnit::Mega), "1.0MB");
165    }
166}
anti_fs/du.rs

anti_fs/
du.rs
