use btlib::Result;
use std::fs::{metadata, read_dir, Metadata};
use std::path::Path;
pub fn visit<P: AsRef<Path>, R, F: FnMut(Metadata) -> R>(path: P, visitor: &mut F) -> Result<()> {
let meta = metadata(&path)?;
if !meta.is_dir() {
visitor(meta);
return Ok(());
}
let contents = read_dir(&path)?;
let mut entry_path = path.as_ref().to_owned();
entry_path.push("x");
for entry in contents {
let entry = entry?;
entry_path.pop();
entry_path.push(entry.file_name());
visit(&entry_path, visitor)?;
}
Ok(())
}
pub fn disk_usage<P: AsRef<Path>>(path: P) -> Result<u64> {
let mut total = 0;
visit(path, &mut |meta| {
if meta.is_file() {
total += meta.len();
}
})?;
Ok(total)
}
pub fn num_files<P: AsRef<Path>>(path: P) -> Result<u64> {
let mut total = 0;
visit(path, &mut |meta| {
if meta.is_file() {
total += 1;
}
})?;
Ok(total)
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs::{create_dir, create_dir_all, write};
use tempdir::TempDir;
#[test]
fn disk_usage_direct_descendants() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3 * BUF.len() as u64;
let dir = TempDir::new("disk_usage").unwrap();
let dir_path = dir.path();
for k in 0..3 {
write(dir_path.join(k.to_string()), &BUF).unwrap();
}
let actual = disk_usage(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn disk_usage_first_gen_dirs() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3 * BUF.len() as u64;
let dir = TempDir::new("disk_usage").unwrap();
let dir_path = dir.path();
write(dir_path.join("1"), &BUF).unwrap();
let sub1 = dir_path.join("sub1");
create_dir(&sub1).unwrap();
write(sub1.join("2"), &BUF).unwrap();
let sub2 = dir_path.join("sub2");
create_dir(&sub2).unwrap();
write(sub2.join("3"), &BUF).unwrap();
let actual = disk_usage(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn disk_usage_second_gen_dir() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3 * BUF.len() as u64;
let dir = TempDir::new("disk_usage").unwrap();
let dir_path = dir.path();
for k in 0..2 {
write(dir_path.join(k.to_string()), &BUF).unwrap();
}
let mut sub = dir_path.to_owned();
sub.push("sub");
sub.push("sub");
create_dir_all(&sub).unwrap();
sub.push("2");
write(&sub, &BUF).unwrap();
let actual = disk_usage(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn disk_usage_empty_dir() {
const EXPECTED: u64 = 0;
let dir = TempDir::new("disk_usage").unwrap();
let dir_path = dir.path();
let actual = disk_usage(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn num_files_direct_descendants() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3;
let dir = TempDir::new("num_files").unwrap();
let dir_path = dir.path();
for k in 0..3 {
write(dir_path.join(k.to_string()), &BUF).unwrap();
}
let actual = num_files(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn num_files_first_gen_dirs() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3;
let dir = TempDir::new("num_files").unwrap();
let dir_path = dir.path();
write(dir_path.join("1"), &BUF).unwrap();
let sub1 = dir_path.join("sub1");
create_dir(&sub1).unwrap();
write(sub1.join("2"), &BUF).unwrap();
let sub2 = dir_path.join("sub2");
create_dir(&sub2).unwrap();
write(sub2.join("3"), &BUF).unwrap();
let actual = num_files(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn num_files_second_gen_dir() {
const BUF: [u8; 8] = [1u8; 8];
const EXPECTED: u64 = 3;
let dir = TempDir::new("num_files").unwrap();
let dir_path = dir.path();
for k in 0..2 {
write(dir_path.join(k.to_string()), &BUF).unwrap();
}
let mut sub = dir_path.to_owned();
sub.push("sub");
sub.push("sub");
create_dir_all(&sub).unwrap();
sub.push("2");
write(&sub, &BUF).unwrap();
let actual = num_files(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
#[test]
fn num_files_empty_dir() {
const EXPECTED: u64 = 0;
let dir = TempDir::new("num_files").unwrap();
let dir_path = dir.path();
let actual = num_files(dir_path).unwrap();
assert_eq!(EXPECTED, actual);
}
}