compare_dir/
progress.rs

1use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
2use indicatif_log_bridge::LogWrapper;
3use std::io::{IsTerminal, stderr};
4use std::path::Path;
5use std::sync::atomic::{self, AtomicBool};
6use std::sync::{Arc, Mutex, Weak};
7use std::time::Duration;
8
9const SPINNER_STYLE0: &str = "{elapsed_precise} {spinner:.green} ";
10const SPINNER_STYLE1_NUM: &str = "{pos:>7} {msg}";
11const SPINNER_STYLE1_SIZE: &str = "{bytes:>7} {msg}";
12const NORMAL_STYLE0: &str = "{elapsed_precise} +{eta:>3} {percent:>3}% {bar:40.cyan/blue} ";
13const NORMAL_STYLE1_NUM: &str = "{pos:>7}/{len:7} {msg}";
14const NORMAL_STYLE1_SIZE: &str = "{bytes:>7}/{total_bytes:7} {msg}";
15const FILE_STYLE: &str = "  {elapsed:>3} +{eta:>3} {percent:>3}% {bar:10.cyan/blue} {wide_msg}";
16
17#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
18pub(crate) struct ProgressValue {
19    pub(crate) num_files: u64,
20    size: u64,
21}
22
23impl ProgressValue {
24    pub(crate) fn with_file_and_size(num_files: u64, size: u64) -> Self {
25        Self { size, num_files }
26    }
27
28    pub(crate) fn with_size(size: u64) -> Self {
29        Self::with_file_and_size(1, size)
30    }
31
32    pub(crate) fn with_skip(size: u64) -> Self {
33        Self::with_file_and_size(1, size)
34    }
35}
36
37impl std::ops::Add for ProgressValue {
38    type Output = Self;
39
40    fn add(self, other: Self) -> Self {
41        Self {
42            num_files: self.num_files + other.num_files,
43            size: self.size + other.size,
44        }
45    }
46}
47
48impl std::ops::AddAssign for ProgressValue {
49    fn add_assign(&mut self, other: Self) {
50        self.num_files += other.num_files;
51        self.size += other.size;
52    }
53}
54
55#[derive(Debug, Default)]
56pub(crate) struct Progress {
57    inner: Option<ProgressBar>,
58    pos: ProgressValue,
59    use_bytes: bool,
60    len: Option<ProgressValue>,
61    multi: Option<MultiProgress>,
62    primary: Weak<Mutex<Progress>>,
63}
64
65impl Progress {
66    pub fn none() -> Self {
67        Self::none_with_primary(Weak::new())
68    }
69
70    fn none_with_primary(primary: Weak<Mutex<Progress>>) -> Self {
71        Self {
72            inner: None,
73            multi: None,
74            primary,
75            ..Default::default()
76        }
77    }
78
79    fn update_style(&self) {
80        if let Some(inner) = &self.inner {
81            let style = if self.len.is_some() {
82                if self.use_bytes {
83                    format!("{NORMAL_STYLE0}{NORMAL_STYLE1_SIZE}")
84                } else {
85                    format!("{NORMAL_STYLE0}{NORMAL_STYLE1_NUM}")
86                }
87            } else {
88                if self.use_bytes {
89                    format!("{SPINNER_STYLE0}{SPINNER_STYLE1_SIZE}")
90                } else {
91                    format!("{SPINNER_STYLE0}{SPINNER_STYLE1_NUM}")
92                }
93            };
94            inner.set_style(ProgressStyle::with_template(&style).unwrap());
95        }
96    }
97
98    fn update_position(&self) {
99        if let Some(inner) = &self.inner {
100            inner.set_position(if self.use_bytes {
101                self.pos.size
102            } else {
103                self.pos.num_files
104            });
105        }
106    }
107
108    fn update_length(&self) {
109        if let Some(inner) = &self.inner
110            && let Some(len) = self.len
111        {
112            inner.set_length(if self.use_bytes {
113                len.size
114            } else {
115                len.num_files
116            });
117        }
118    }
119
120    pub fn set_message(&self, msg: impl Into<String>) {
121        if let Some(inner) = &self.inner {
122            inner.set_message(msg.into());
123        }
124    }
125
126    pub fn inc(&mut self, amount: ProgressValue) {
127        self.pos += amount;
128        self.update_position();
129        if let Some(primary_arc) = self.primary.upgrade() {
130            primary_arc.lock().unwrap().inc(amount);
131        }
132    }
133
134    pub fn inc_size(&mut self, size: u64) {
135        let amount = ProgressValue::with_file_and_size(0, size);
136        self.inc(amount);
137    }
138
139    pub fn inc_file(&mut self, num_files: u64) {
140        let amount = ProgressValue::with_file_and_size(num_files, 0);
141        self.inc(amount);
142    }
143
144    pub fn set_length(&mut self, len: ProgressValue) {
145        self.len = Some(len);
146        self.update_style();
147        self.update_length();
148    }
149
150    pub fn use_bytes(&mut self) {
151        self.use_bytes = true;
152        self.update_style();
153        self.update_length();
154        self.update_position();
155    }
156
157    pub fn finish(&self) {
158        if let Some(len) = self.len {
159            assert_eq!(self.pos, len);
160        }
161        if let Some(inner) = &self.inner {
162            inner.finish();
163            if let Some(multi) = &self.multi {
164                multi.remove(inner);
165            }
166        }
167    }
168
169    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
170    where
171        F: FnOnce() -> R,
172    {
173        if !stream.is_terminal() {
174            return f();
175        }
176        if let Some(multi) = &self.multi {
177            multi.suspend(f)
178        } else if let Some(inner) = &self.inner {
179            inner.suspend(f)
180        } else {
181            f()
182        }
183    }
184}
185
186#[derive(Clone)]
187pub(crate) struct SharedProgress {
188    inner: Arc<Mutex<Progress>>,
189}
190
191impl SharedProgress {
192    pub fn none() -> Self {
193        Self {
194            inner: Arc::new(Mutex::new(Progress::none())),
195        }
196    }
197
198    pub fn set_message(&self, msg: impl Into<String>) {
199        self.inner.lock().unwrap().set_message(msg);
200    }
201
202    pub fn inc(&self, amount: ProgressValue) {
203        self.inner.lock().unwrap().inc(amount);
204    }
205
206    pub fn set_length(&self, len: ProgressValue) {
207        self.inner.lock().unwrap().set_length(len);
208    }
209
210    pub fn use_bytes(&self) {
211        self.inner.lock().unwrap().use_bytes();
212    }
213
214    pub fn finish(&self) {
215        self.inner.lock().unwrap().finish();
216    }
217
218    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
219    where
220        F: FnOnce() -> R,
221    {
222        self.inner.lock().unwrap().suspend_for(stream, f)
223    }
224}
225
226#[derive(Debug)]
227pub struct ProgressBuilder {
228    multi: MultiProgress,
229    pub is_enabled: bool,
230    pub is_file_enabled: bool,
231    is_propagate: AtomicBool,
232    primary: Mutex<Weak<Mutex<Progress>>>,
233}
234
235impl Default for ProgressBuilder {
236    fn default() -> Self {
237        Self {
238            multi: MultiProgress::default(),
239            is_enabled: stderr().is_terminal(),
240            is_file_enabled: false,
241            is_propagate: AtomicBool::new(false),
242            primary: Mutex::new(Weak::new()),
243        }
244    }
245}
246
247impl ProgressBuilder {
248    pub fn new() -> Self {
249        Self::default()
250    }
251
252    pub fn init_logger(&self, logger: env_logger::Logger) -> anyhow::Result<()> {
253        let max_level = logger.filter();
254        LogWrapper::new(self.multi.clone(), logger).try_init()?;
255        log::set_max_level(max_level);
256        Ok(())
257    }
258
259    pub(crate) fn is_propagate(&self) -> bool {
260        self.is_propagate.load(atomic::Ordering::Relaxed)
261    }
262
263    pub(crate) fn set_propagate(&self) {
264        self.is_propagate.store(true, atomic::Ordering::Relaxed);
265    }
266
267    pub(crate) fn add_primary(&self) -> SharedProgress {
268        if !self.is_enabled {
269            return SharedProgress::none();
270        }
271        let inner = self.multi.add(ProgressBar::new_spinner());
272        inner.enable_steady_tick(Duration::from_secs(1));
273        let progress = Progress {
274            inner: Some(inner),
275            multi: Some(self.multi.clone()),
276            primary: Weak::new(),
277            ..Default::default()
278        };
279        progress.update_style();
280        let shared = SharedProgress {
281            inner: Arc::new(Mutex::new(progress)),
282        };
283        *self.primary.lock().unwrap() = Arc::downgrade(&shared.inner);
284        shared
285    }
286
287    pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
288        if !self.is_enabled {
289            return Progress::none();
290        }
291        let primary = if self.is_propagate() {
292            Weak::clone(&self.primary.lock().unwrap())
293        } else {
294            Weak::new()
295        };
296        if !self.is_file_enabled {
297            return Progress::none_with_primary(primary);
298        }
299
300        let inner = self.multi.add(ProgressBar::new(file_size));
301        inner.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
302        if let Some(parent) = path.parent()
303            && let Some(file_name) = path.file_name()
304        {
305            inner.set_message(format!(
306                "{} ({})",
307                file_name.to_string_lossy(),
308                parent.to_string_lossy()
309            ));
310        } else {
311            inner.set_message(path.to_string_lossy().to_string());
312        }
313        Progress {
314            inner: Some(inner),
315            multi: Some(self.multi.clone()),
316            use_bytes: true,
317            primary,
318            ..Default::default()
319        }
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326
327    #[test]
328    fn propagate() {
329        let builder = ProgressBuilder {
330            is_enabled: true,
331            is_file_enabled: false,
332            ..Default::default()
333        };
334        builder.set_propagate();
335        let _primary = builder.add_primary();
336
337        // Add a file progress
338        let file_path = Path::new("dummy.txt");
339        let mut file_progress = builder.add_file(file_path, 100);
340
341        // Initially primary position is 0
342        {
343            let prim = builder.primary.lock().unwrap().upgrade().unwrap();
344            assert_eq!(prim.lock().unwrap().pos.size, 0);
345        }
346
347        // Increment file progress
348        file_progress.inc(ProgressValue::with_size(10));
349
350        // Primary progress should be incremented to 10
351        {
352            let prim = builder.primary.lock().unwrap().upgrade().unwrap();
353            assert_eq!(prim.lock().unwrap().pos.size, 10);
354            assert_eq!(prim.lock().unwrap().pos.num_files, 1);
355        }
356    }
357}
compare_dir/progress.rs

compare_dir/
progress.rs
