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<F, R>(&self, f: F) -> R
60    where
61        F: FnOnce() -> R,
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: true,
86        }
87    }
88}
89
90impl ProgressBuilder {
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    pub(crate) fn add_spinner(&self) -> Progress {
96        if !self.is_enabled {
97            return Progress::none();
98        }
99        let progress = self.multi.add(ProgressBar::new_spinner());
100        progress.enable_steady_tick(Duration::from_millis(120));
101        progress.set_style(ProgressStyle::with_template(SPINNER_STYLE).unwrap());
102        Progress {
103            inner: Some(progress),
104            multi: Some(self.multi.clone()),
105        }
106    }
107
108    pub(crate) fn add_file(&self, path: &Path, file_size: u64) -> Progress {
109        if !self.is_enabled || !self.is_file_enabled {
110            return Progress::none();
111        }
112        let progress = self.multi.add(ProgressBar::new(file_size));
113        progress.set_style(ProgressStyle::with_template(FILE_STYLE).unwrap());
114        if let Some(parent) = path.parent()
115            && let Some(file_name) = path.file_name()
116        {
117            progress.set_message(format!(
118                "{} ({})",
119                file_name.to_string_lossy(),
120                parent.to_string_lossy()
121            ));
122        } else {
123            progress.set_message(path.to_string_lossy().to_string());
124        }
125        Progress {
126            inner: Some(progress),
127            multi: Some(self.multi.clone()),
128        }
129    }
130}