dua/
aggregate.rs

1use crate::{crossdev, ByteFormat, InodeFilter, Throttle, WalkOptions, WalkResult};
2use anyhow::Result;
3use filesize::PathExt;
4use owo_colors::{AnsiColors as Color, OwoColorize};
5use std::time::Duration;
6use std::{io, path::Path};
7
8/// Aggregate the given `paths` and write information about them to `out` in a human-readable format.
9/// If `compute_total` is set, it will write an additional line with the total size across all given `paths`.
10/// If `sort_by_size_in_bytes` is set, we will sort all sizes (ascending) before outputting them.
11pub fn aggregate(
12    mut out: impl io::Write,
13    mut err: Option<impl io::Write>,
14    walk_options: WalkOptions,
15    compute_total: bool,
16    sort_by_size_in_bytes: bool,
17    byte_format: ByteFormat,
18    paths: impl IntoIterator<Item = impl AsRef<Path>>,
19) -> Result<(WalkResult, Statistics)> {
20    let mut res = WalkResult::default();
21    let mut stats = Statistics {
22        smallest_file_in_bytes: u128::MAX,
23        ..Default::default()
24    };
25    let mut total = 0;
26    let mut num_roots = 0;
27    let mut aggregates = Vec::new();
28    let mut inodes = InodeFilter::default();
29    let progress = Throttle::new(Duration::from_millis(100), Duration::from_secs(1).into());
30
31    for path in paths.into_iter() {
32        num_roots += 1;
33        let mut num_bytes = 0u128;
34        let mut num_errors = 0u64;
35        let device_id = match crossdev::init(path.as_ref()) {
36            Ok(id) => id,
37            Err(_) => {
38                num_errors += 1;
39                res.num_errors += 1;
40                aggregates.push((path.as_ref().to_owned(), num_bytes, num_errors));
41                continue;
42            }
43        };
44        for entry in walk_options.iter_from_path(path.as_ref(), device_id, false) {
45            stats.entries_traversed += 1;
46            progress.throttled(|| {
47                if let Some(err) = err.as_mut() {
48                    write!(err, "Enumerating {} items\r", stats.entries_traversed).ok();
49                }
50            });
51            match entry {
52                Ok(entry) => {
53                    let file_size = match entry.client_state {
54                        Some(Ok(ref m))
55                            if !m.is_dir()
56                                && (walk_options.count_hard_links || inodes.add(m))
57                                && (walk_options.cross_filesystems
58                                    || crossdev::is_same_device(device_id, m)) =>
59                        {
60                            if walk_options.apparent_size {
61                                m.len()
62                            } else {
63                                entry.path().size_on_disk_fast(m).unwrap_or_else(|_| {
64                                    num_errors += 1;
65                                    0
66                                })
67                            }
68                        }
69                        Some(Ok(_)) => 0,
70                        Some(Err(_)) => {
71                            num_errors += 1;
72                            0
73                        }
74                        None => 0, // ignore directory
75                    } as u128;
76                    stats.largest_file_in_bytes = stats.largest_file_in_bytes.max(file_size);
77                    stats.smallest_file_in_bytes = stats.smallest_file_in_bytes.min(file_size);
78                    num_bytes += file_size;
79                }
80                Err(_) => num_errors += 1,
81            }
82        }
83
84        if let Some(err) = err.as_mut() {
85            write!(err, "\x1b[2K\r").ok();
86        }
87
88        if sort_by_size_in_bytes {
89            aggregates.push((path.as_ref().to_owned(), num_bytes, num_errors));
90        } else {
91            output_colored_path(
92                &mut out,
93                &path,
94                num_bytes,
95                num_errors,
96                path_color_of(&path),
97                byte_format,
98            )?;
99        }
100        total += num_bytes;
101        res.num_errors += num_errors;
102    }
103
104    if stats.entries_traversed == 0 {
105        stats.smallest_file_in_bytes = 0;
106    }
107
108    if sort_by_size_in_bytes {
109        aggregates.sort_by_key(|&(_, num_bytes, _)| num_bytes);
110        for (path, num_bytes, num_errors) in aggregates.into_iter() {
111            output_colored_path(
112                &mut out,
113                &path,
114                num_bytes,
115                num_errors,
116                path_color_of(&path),
117                byte_format,
118            )?;
119        }
120    }
121
122    if num_roots > 1 && compute_total {
123        output_colored_path(
124            &mut out,
125            Path::new("total"),
126            total,
127            res.num_errors,
128            None,
129            byte_format,
130        )?;
131    }
132    Ok((res, stats))
133}
134
135fn path_color_of(path: impl AsRef<Path>) -> Option<Color> {
136    (!path.as_ref().is_file()).then_some(Color::Cyan)
137}
138
139fn output_colored_path(
140    out: &mut impl io::Write,
141    path: impl AsRef<Path>,
142    num_bytes: u128,
143    num_errors: u64,
144    path_color: Option<Color>,
145    byte_format: ByteFormat,
146) -> std::result::Result<(), io::Error> {
147    let size = byte_format.display(num_bytes).to_string();
148    let size = size.green();
149    let size_width = byte_format.width();
150    let path = path.as_ref().display();
151
152    let errors = (num_errors != 0)
153        .then(|| {
154            let plural_s = if num_errors > 1 { "s" } else { "" };
155            format!("  <{num_errors} IO Error{plural_s}>")
156        })
157        .unwrap_or_default();
158
159    if let Some(color) = path_color {
160        writeln!(out, "{size:>size_width$} {}{errors}", path.color(color))
161    } else {
162        writeln!(out, "{size:>size_width$} {path}{errors}")
163    }
164}
165
166/// Statistics obtained during a filesystem walk
167#[derive(Default, Debug)]
168pub struct Statistics {
169    /// The amount of entries we have seen during filesystem traversal
170    pub entries_traversed: u64,
171    /// The size of the smallest file encountered in bytes
172    pub smallest_file_in_bytes: u128,
173    /// The size of the largest file encountered in bytes
174    pub largest_file_in_bytes: u128,
175}