dir_size/
lib.rs

1// SPDX-FileCopyrightText: 2024 Integral <integral@member.fsf.org>
2//
3// SPDX-License-Identifier: MPL-2.0
4
5use rayon::prelude::*;
6use std::{fs, io, path::Path};
7
8const KIBIBYTE: u64 = 1 << 10;
9const MEBIBYTE: u64 = 1 << 20;
10const GIBIBYTE: u64 = 1 << 30;
11const TEBIBYTE: u64 = 1 << 40;
12const PEBIBYTE: u64 = 1 << 50;
13const EXBIBYTE: u64 = 1 << 60;
14
15/// Get the size of the file (in bytes).
16///
17/// If `path` points to a directory, calculate the size of directory recursively,
18/// including all of its files and subdirectories.
19///
20/// This function will return an error if `path` does not exist,
21/// or user lacks permissions to perform `metadata` call on `path`.
22pub fn get_size_in_bytes(path: &Path) -> io::Result<u64> {
23    match fs::symlink_metadata(path) {
24        Ok(meta) if meta.is_file() => Ok(meta.len()),
25        Ok(meta) if meta.is_dir() => get_dir_size(path),
26        Ok(_) => Ok(0),
27        Err(e) => Err(e),
28    }
29}
30
31/// Get the size of the file (in human-readable bytes).
32///
33/// If `path` points to a directory, calculate the size of directory recursively,
34/// including all of its files and subdirectories.
35///
36/// This function will return an error if `path` does not exist,
37/// or user lacks permissions to perform `metadata` call on `path`.
38pub fn get_size_in_human_bytes(path: &Path) -> io::Result<String> {
39    Ok(convert_to_human_bytes(get_size_in_bytes(path)?, false))
40}
41
42/// Get the size of the file (in human-readable bytes, using abbreviated units (K, M, G, etc.))
43///
44/// If `path` points to a directory, calculate the size of directory recursively,
45/// including all of its files and subdirectories.
46///
47/// This function will return an error if `path` does not exist,
48/// or user lacks permissions to perform `metadata` call on `path`.
49pub fn get_size_in_abbr_human_bytes(path: &Path) -> io::Result<String> {
50    Ok(convert_to_human_bytes(get_size_in_bytes(path)?, true))
51}
52
53fn get_dir_size(path: &Path) -> io::Result<u64> {
54    let entries: Vec<_> = fs::read_dir(path)?.collect();
55
56    let total = entries
57        .par_iter()
58        .filter_map(|entry| match entry {
59            Ok(entry) => match entry.metadata() {
60                Ok(meta) if meta.is_file() => Some(meta.len()),
61                Ok(meta) if meta.is_dir() => get_dir_size(&entry.path()).ok(),
62                _ => None,
63            },
64            _ => None,
65        })
66        .sum();
67
68    Ok(total)
69}
70
71fn convert_to_human_bytes(size_in_bytes: u64, abbr: bool) -> String {
72    for ((min_bytes, max_bytes), abbr_unit, full_unit) in [
73        ((1, KIBIBYTE), "B", "Bytes"),      // Bytes
74        ((KIBIBYTE, MEBIBYTE), "K", "KiB"), // KiB
75        ((MEBIBYTE, GIBIBYTE), "M", "MiB"), // MiB
76        ((GIBIBYTE, TEBIBYTE), "G", "GiB"), // GiB
77        ((TEBIBYTE, PEBIBYTE), "T", "TiB"), // TiB
78        ((PEBIBYTE, EXBIBYTE), "P", "PiB"), // PiB
79    ] {
80        if size_in_bytes < max_bytes {
81            return format!(
82                "{} {}",
83                size_in_bytes / min_bytes,
84                if abbr { abbr_unit } else { full_unit }
85            );
86        }
87    }
88
89    format!(
90        "{} {}",
91        size_in_bytes / EXBIBYTE,
92        if abbr { "E" } else { "EiB" }
93    )
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn test_convert_to_human_bytes() {
102        for (size_in_bytes, human_bytes) in [
103            (0, "0 Bytes"),
104            (KIBIBYTE - 1, "1023 Bytes"),
105            (KIBIBYTE, "1 KiB"),
106            (MEBIBYTE - 1, "1023 KiB"),
107            (MEBIBYTE, "1 MiB"),
108            (GIBIBYTE - 1, "1023 MiB"),
109            (GIBIBYTE, "1 GiB"),
110            (TEBIBYTE - 1, "1023 GiB"),
111            (TEBIBYTE, "1 TiB"),
112            (PEBIBYTE - 1, "1023 TiB"),
113            (PEBIBYTE, "1 PiB"),
114            (EXBIBYTE - 1, "1023 PiB"),
115            (EXBIBYTE, "1 EiB"),
116        ] {
117            println!("{size_in_bytes} bytes -> {human_bytes}");
118            assert_eq!(convert_to_human_bytes(size_in_bytes, false), human_bytes);
119        }
120    }
121
122    #[test]
123    fn test_convert_to_abbr_human_bytes() {
124        for (size_in_bytes, abbr_human_bytes) in [
125            (0, "0 B"),
126            (KIBIBYTE - 1, "1023 B"),
127            (KIBIBYTE, "1 K"),
128            (MEBIBYTE - 1, "1023 K"),
129            (MEBIBYTE, "1 M"),
130            (GIBIBYTE - 1, "1023 M"),
131            (GIBIBYTE, "1 G"),
132            (TEBIBYTE - 1, "1023 G"),
133            (TEBIBYTE, "1 T"),
134            (PEBIBYTE - 1, "1023 T"),
135            (PEBIBYTE, "1 P"),
136            (EXBIBYTE - 1, "1023 P"),
137            (EXBIBYTE, "1 E"),
138        ] {
139            println!("{size_in_bytes} bytes -> {abbr_human_bytes}");
140            assert_eq!(
141                convert_to_human_bytes(size_in_bytes, true),
142                abbr_human_bytes
143            );
144        }
145    }
146}
dir_size/lib.rs

dir_size/
lib.rs
