Skip to main content

compare_dir/
progress.rs

1use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
2use std::io::{IsTerminal, stderr};
3use std::path::Path;
4use std::time::Duration;
5
6const SPINNER_STYLE: &str = "{elapsed_precise} {spinner:.green} {pos:>7} {msg}";
7const NORMAL_STYLE: &str =
8    "{elapsed_precise} +{eta:>3} {percent:>3}% {bar:40.cyan/blue} {pos:>7}/{len:7} {msg}";
9const FILE_STYLE: &str = "  {elapsed:>3} +{eta:>3} {percent:>3}% {wide_msg}";
10
11#[derive(Debug)]
12pub(crate) struct Progress {
13    inner: Option<ProgressBar>,
14    multi: Option<MultiProgress>,
15}
16
17impl Progress {
18    pub fn none() -> Self {
19        Self {
20            inner: None,
21            multi: None,
22        }
23    }
24
25    pub fn set_message(&self, msg: impl Into<String>) {
26        if let Some(inner) = &self.inner {
27            inner.set_message(msg.into());
28        }
29    }
30
31    pub fn inc(&self, amount: u64) {
32        if let Some(inner) = &self.inner {
33            inner.inc(amount);
34        }
35    }
36
37    pub fn length(&self) -> Option<u64> {
38        self.inner.as_ref().and_then(|inner| inner.length())
39    }
40
41    pub fn set_length(&self, len: u64) {
42        if let Some(inner) = &self.inner {
43            if inner.length().is_none() {
44                inner.set_style(ProgressStyle::with_template(NORMAL_STYLE).unwrap());
45            }
46            inner.set_length(len);
47        }
48    }
49
50    pub fn finish(&self) {
51        if let Some(inner) = &self.inner {
52            inner.finish();
53            if let Some(multi) = &self.multi {
54                multi.remove(inner);
55            }
56        }
57    }
58
59    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
60    where
61        F: FnOnce() -> R,
62    {
63        if !stream.is_terminal() {
64            return f();
65        }
66        if let Some(multi) = &self.multi {
67            multi.suspend(f)
68        } else if let Some(inner) = &self.inner {
69            inner.suspend(f)
70        } else {
71            f()
72        }
73    }
74}
75
76#[derive(Debug)]
77pub struct ProgressBuilder {
78    multi: MultiProgress,
79    pub is_enabled: bool,
80    pub is_file_enabled: bool,
81}
82
83impl Default for ProgressBuilder {
84    fn default() -> Self {
85        Self {
86            multi: MultiProgress::default(),
87            is_enabled: stderr().is_terminal(),
88            is_file_enabled: false,
89        }
90    }
91}
92
93impl ProgressBuilder {
94    pub fn new() -> Self {
95        Self::default()
96    }
97
98    pub(crate) fn add_spinner(&self) -> Progress {
99        if !self.is_enabled {
100            return Progress::none();
101        }
102        let progress = self.multi.add(ProgressBar::new_spinner());
103        progress.enable_steady_tick(Duration::from_secs(1));
104        progress.set_style(ProgressStyle::with_template(SPINNER_STYLE).unwrap());
105        Progress {
106            inner: Some(progress),
107            multi: Some(self.multi.clone()),
108        }
109    }
110
111    pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
112        if !self.is_enabled || !self.is_file_enabled {
113            return Progress::none();
114        }
115        let progress = self.multi.add(ProgressBar::new(file_size));
116        progress.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
117        if let Some(parent) = path.parent()
118            && let Some(file_name) = path.file_name()
119        {
120            progress.set_message(format!(
121                "{} ({})",
122                file_name.to_string_lossy(),
123                parent.to_string_lossy()
124            ));
125        } else {
126            progress.set_message(path.to_string_lossy().to_string());
127        }
128        Progress {
129            inner: Some(progress),
130            multi: Some(self.multi.clone()),
131        }
132    }
133}