Skip to main content

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