fs_queries.rs 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  1. use btlib::Result;
  2. use std::fs::{metadata, read_dir, Metadata};
  3. use std::path::Path;
  4. /// Calls the given closure on every entry recursively contained in the given path.
  5. pub fn visit<P: AsRef<Path>, R, F: FnMut(Metadata) -> R>(path: P, visitor: &mut F) -> Result<()> {
  6. let meta = metadata(&path)?;
  7. if !meta.is_dir() {
  8. visitor(meta);
  9. return Ok(());
  10. }
  11. let contents = read_dir(&path)?;
  12. let mut entry_path = path.as_ref().to_owned();
  13. entry_path.push("x");
  14. for entry in contents {
  15. let entry = entry?;
  16. entry_path.pop();
  17. entry_path.push(entry.file_name());
  18. visit(&entry_path, visitor)?;
  19. }
  20. Ok(())
  21. }
  22. /// Recursively sums up the length of each file stored under the given path and returns the total.
  23. /// Note that if the given path contains hard links, the disk usage for the hard linked file will
  24. /// be counted multiple times, once for each link.
  25. pub fn disk_usage<P: AsRef<Path>>(path: P) -> Result<u64> {
  26. let mut total = 0;
  27. visit(path, &mut |meta| {
  28. if meta.is_file() {
  29. total += meta.len();
  30. }
  31. })?;
  32. Ok(total)
  33. }
  34. pub fn num_files<P: AsRef<Path>>(path: P) -> Result<u64> {
  35. let mut total = 0;
  36. visit(path, &mut |meta| {
  37. if meta.is_file() {
  38. total += 1;
  39. }
  40. })?;
  41. Ok(total)
  42. }
  43. #[cfg(test)]
  44. mod tests {
  45. use super::*;
  46. use std::fs::{create_dir, create_dir_all, write};
  47. use tempdir::TempDir;
  48. #[test]
  49. fn disk_usage_direct_descendants() {
  50. const BUF: [u8; 8] = [1u8; 8];
  51. const EXPECTED: u64 = 3 * BUF.len() as u64;
  52. let dir = TempDir::new("disk_usage").unwrap();
  53. let dir_path = dir.path();
  54. for k in 0..3 {
  55. write(dir_path.join(k.to_string()), &BUF).unwrap();
  56. }
  57. let actual = disk_usage(dir_path).unwrap();
  58. assert_eq!(EXPECTED, actual);
  59. }
  60. #[test]
  61. fn disk_usage_first_gen_dirs() {
  62. const BUF: [u8; 8] = [1u8; 8];
  63. const EXPECTED: u64 = 3 * BUF.len() as u64;
  64. let dir = TempDir::new("disk_usage").unwrap();
  65. let dir_path = dir.path();
  66. write(dir_path.join("1"), &BUF).unwrap();
  67. let sub1 = dir_path.join("sub1");
  68. create_dir(&sub1).unwrap();
  69. write(sub1.join("2"), &BUF).unwrap();
  70. let sub2 = dir_path.join("sub2");
  71. create_dir(&sub2).unwrap();
  72. write(sub2.join("3"), &BUF).unwrap();
  73. let actual = disk_usage(dir_path).unwrap();
  74. assert_eq!(EXPECTED, actual);
  75. }
  76. #[test]
  77. fn disk_usage_second_gen_dir() {
  78. const BUF: [u8; 8] = [1u8; 8];
  79. const EXPECTED: u64 = 3 * BUF.len() as u64;
  80. let dir = TempDir::new("disk_usage").unwrap();
  81. let dir_path = dir.path();
  82. for k in 0..2 {
  83. write(dir_path.join(k.to_string()), &BUF).unwrap();
  84. }
  85. let mut sub = dir_path.to_owned();
  86. sub.push("sub");
  87. sub.push("sub");
  88. create_dir_all(&sub).unwrap();
  89. sub.push("2");
  90. write(&sub, &BUF).unwrap();
  91. let actual = disk_usage(dir_path).unwrap();
  92. assert_eq!(EXPECTED, actual);
  93. }
  94. #[test]
  95. fn disk_usage_empty_dir() {
  96. const EXPECTED: u64 = 0;
  97. let dir = TempDir::new("disk_usage").unwrap();
  98. let dir_path = dir.path();
  99. let actual = disk_usage(dir_path).unwrap();
  100. assert_eq!(EXPECTED, actual);
  101. }
  102. #[test]
  103. fn num_files_direct_descendants() {
  104. const BUF: [u8; 8] = [1u8; 8];
  105. const EXPECTED: u64 = 3;
  106. let dir = TempDir::new("num_files").unwrap();
  107. let dir_path = dir.path();
  108. for k in 0..3 {
  109. write(dir_path.join(k.to_string()), &BUF).unwrap();
  110. }
  111. let actual = num_files(dir_path).unwrap();
  112. assert_eq!(EXPECTED, actual);
  113. }
  114. #[test]
  115. fn num_files_first_gen_dirs() {
  116. const BUF: [u8; 8] = [1u8; 8];
  117. const EXPECTED: u64 = 3;
  118. let dir = TempDir::new("num_files").unwrap();
  119. let dir_path = dir.path();
  120. write(dir_path.join("1"), &BUF).unwrap();
  121. let sub1 = dir_path.join("sub1");
  122. create_dir(&sub1).unwrap();
  123. write(sub1.join("2"), &BUF).unwrap();
  124. let sub2 = dir_path.join("sub2");
  125. create_dir(&sub2).unwrap();
  126. write(sub2.join("3"), &BUF).unwrap();
  127. let actual = num_files(dir_path).unwrap();
  128. assert_eq!(EXPECTED, actual);
  129. }
  130. #[test]
  131. fn num_files_second_gen_dir() {
  132. const BUF: [u8; 8] = [1u8; 8];
  133. const EXPECTED: u64 = 3;
  134. let dir = TempDir::new("num_files").unwrap();
  135. let dir_path = dir.path();
  136. for k in 0..2 {
  137. write(dir_path.join(k.to_string()), &BUF).unwrap();
  138. }
  139. let mut sub = dir_path.to_owned();
  140. sub.push("sub");
  141. sub.push("sub");
  142. create_dir_all(&sub).unwrap();
  143. sub.push("2");
  144. write(&sub, &BUF).unwrap();
  145. let actual = num_files(dir_path).unwrap();
  146. assert_eq!(EXPECTED, actual);
  147. }
  148. #[test]
  149. fn num_files_empty_dir() {
  150. const EXPECTED: u64 = 0;
  151. let dir = TempDir::new("num_files").unwrap();
  152. let dir_path = dir.path();
  153. let actual = num_files(dir_path).unwrap();
  154. assert_eq!(EXPECTED, actual);
  155. }
  156. }