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 (walk_options.count_hard_links || inodes.add(m))
56                                && (walk_options.cross_filesystems
57                                    || crossdev::is_same_device(device_id, m)) =>
58                        {
59                            if walk_options.apparent_size {
60                                m.len()
61                            } else {
62                                entry.path().size_on_disk_fast(m).unwrap_or_else(|_| {
63                                    num_errors += 1;
64                                    0
65                                })
66                            }
67                        }
68                        Some(Ok(_)) => 0,
69                        Some(Err(_)) => {
70                            num_errors += 1;
71                            0
72                        }
73                        None => 0, // ignore directory
74                    } as u128;
75                    stats.largest_file_in_bytes = stats.largest_file_in_bytes.max(file_size);
76                    stats.smallest_file_in_bytes = stats.smallest_file_in_bytes.min(file_size);
77                    num_bytes += file_size;
78                }
79                Err(_) => num_errors += 1,
80            }
81        }
82
83        if let Some(err) = err.as_mut() {
84            write!(err, "\x1b[2K\r").ok();
85        }
86
87        if sort_by_size_in_bytes {
88            aggregates.push((path.as_ref().to_owned(), num_bytes, num_errors));
89        } else {
90            output_colored_path(
91                &mut out,
92                &path,
93                num_bytes,
94                num_errors,
95                path_color_of(&path),
96                byte_format,
97            )?;
98        }
99        total += num_bytes;
100        res.num_errors += num_errors;
101    }
102
103    if stats.entries_traversed == 0 {
104        stats.smallest_file_in_bytes = 0;
105    }
106
107    if sort_by_size_in_bytes {
108        aggregates.sort_by_key(|&(_, num_bytes, _)| num_bytes);
109        for (path, num_bytes, num_errors) in aggregates.into_iter() {
110            output_colored_path(
111                &mut out,
112                &path,
113                num_bytes,
114                num_errors,
115                path_color_of(&path),
116                byte_format,
117            )?;
118        }
119    }
120
121    if num_roots > 1 && compute_total {
122        output_colored_path(
123            &mut out,
124            Path::new("total"),
125            total,
126            res.num_errors,
127            None,
128            byte_format,
129        )?;
130    }
131    Ok((res, stats))
132}
133
134fn path_color_of(path: impl AsRef<Path>) -> Option<Color> {
135    (!path.as_ref().is_file()).then_some(Color::Cyan)
136}
137
138fn output_colored_path(
139    out: &mut impl io::Write,
140    path: impl AsRef<Path>,
141    num_bytes: u128,
142    num_errors: u64,
143    path_color: Option<Color>,
144    byte_format: ByteFormat,
145) -> std::result::Result<(), io::Error> {
146    let size = byte_format.display(num_bytes).to_string();
147    let size = size.green();
148    let size_width = byte_format.width();
149    let path = path.as_ref().display();
150
151    let errors = if num_errors != 0 {
152        format!(
153            "  <{num_errors} IO Error{plural_s}>",
154            plural_s = if num_errors > 1 { "s" } else { "" }
155        )
156    } else {
157        "".into()
158    };
159
160    if let Some(color) = path_color {
161        writeln!(out, "{size:>size_width$} {}{errors}", path.color(color))
162    } else {
163        writeln!(out, "{size:>size_width$} {path}{errors}")
164    }
165}
166
167/// Statistics obtained during a filesystem walk
168#[derive(Default, Debug)]
169pub struct Statistics {
170    /// The amount of entries we have seen during filesystem traversal
171    pub entries_traversed: u64,
172    /// The size of the smallest file encountered in bytes
173    pub smallest_file_in_bytes: u128,
174    /// The size of the largest file encountered in bytes
175    pub largest_file_in_bytes: u128,
176}