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 set_length(&self, len: u64) {
39        if let Some(inner) = &self.inner {
40            if inner.length().is_none() {
41                inner.set_style(ProgressStyle::with_template(NORMAL_STYLE).unwrap());
42            }
43            inner.set_length(len);
44        }
45    }
46
47    pub fn finish(&self) {
48        if let Some(inner) = &self.inner {
49            inner.finish();
50            if let Some(multi) = &self.multi {
51                multi.remove(inner);
52            }
53        }
54    }
55
56    pub fn suspend_for<F, R, S: IsTerminal>(&self, stream: S, f: F) -> R
57    where
58        F: FnOnce() -> R,
59    {
60        if !stream.is_terminal() {
61            return f();
62        }
63        if let Some(multi) = &self.multi {
64            multi.suspend(f)
65        } else if let Some(inner) = &self.inner {
66            inner.suspend(f)
67        } else {
68            f()
69        }
70    }
71}
72
73#[derive(Debug)]
74pub struct ProgressBuilder {
75    multi: MultiProgress,
76    pub is_enabled: bool,
77    pub is_file_enabled: bool,
78}
79
80impl Default for ProgressBuilder {
81    fn default() -> Self {
82        Self {
83            multi: MultiProgress::default(),
84            is_enabled: stderr().is_terminal(),
85            is_file_enabled: false,
86        }
87    }
88}
89
90impl ProgressBuilder {
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    pub fn init_logger(&self, logger: env_logger::Logger) -> anyhow::Result<()> {
96        let max_level = logger.filter();
97        LogWrapper::new(self.multi.clone(), logger).try_init()?;
98        log::set_max_level(max_level);
99        Ok(())
100    }
101
102    pub(crate) fn add_spinner(&self) -> Progress {
103        if !self.is_enabled {
104            return Progress::none();
105        }
106        let progress = self.multi.add(ProgressBar::new_spinner());
107        progress.enable_steady_tick(Duration::from_secs(1));
108        progress.set_style(ProgressStyle::with_template(SPINNER_STYLE).unwrap());
109        Progress {
110            inner: Some(progress),
111            multi: Some(self.multi.clone()),
112        }
113    }
114
115    pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
116        if !self.is_enabled || !self.is_file_enabled {
117            return Progress::none();
118        }
119        let progress = self.multi.add(ProgressBar::new(file_size));
120        progress.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
121        if let Some(parent) = path.parent()
122            && let Some(file_name) = path.file_name()
123        {
124            progress.set_message(format!(
125                "{} ({})",
126                file_name.to_string_lossy(),
127                parent.to_string_lossy()
128            ));
129        } else {
130            progress.set_message(path.to_string_lossy().to_string());
131        }
132        Progress {
133            inner: Some(progress),
134            multi: Some(self.multi.clone()),
135        }
136    }
137}