watchso/
progress.rs

1//! Progress bars and spinners with consistent behaviour and styles.
2
3use std::future::Future;
4
5use console::Emoji;
6use indicatif::{ProgressBar, ProgressStyle};
7use tokio::time::Duration;
8
9use crate::constants::emoji;
10
11/// Terminal progress utility struct.
12#[derive(Default)]
13pub struct Progress<'a> {
14    message: Option<&'a str>,
15    success_message: Option<&'a str>,
16    error_message: Option<&'a str>,
17    clear: bool,
18}
19
20impl<'a> Progress<'a> {
21    /// Create a new [`Progress`].
22    pub fn new() -> Self {
23        Self::default()
24    }
25
26    /// Set the message that will be displayed while the progress is ongoing.
27    pub fn message(&mut self, message: &'a str) -> &mut Self {
28        self.message = Some(message);
29        self
30    }
31
32    /// Set the success message for the progress.
33    pub fn success_message(&mut self, message: &'a str) -> &mut Self {
34        self.success_message = Some(message);
35        self
36    }
37
38    /// Set the error message for the progress.
39    pub fn error_message(&mut self, message: &'a str) -> &mut Self {
40        self.error_message = Some(message);
41        self
42    }
43
44    /// Set whether the line should be cleared after the progress is finished.
45    #[allow(dead_code)]
46    pub fn clear(&mut self, clear: bool) -> &mut Self {
47        self.clear = clear;
48        self
49    }
50
51    /// Spawn a spinner with the given callback.
52    pub async fn spinner_with<F, R, O>(&self, cb: F) -> miette::Result<O>
53    where
54        F: Fn() -> R,
55        R: Future<Output = miette::Result<O>>,
56    {
57        let pb = ProgressBar::new_spinner();
58        pb.set_style(ProgressStyle::with_template(" {spinner:.green} {msg}").unwrap());
59        pb.enable_steady_tick(Duration::from_millis(120));
60
61        if let Some(message) = self.message {
62            pb.set_message(message.to_owned());
63        }
64
65        let output = cb().await;
66
67        match output {
68            Ok(_) => handle_output(&pb, self.success_message, "green", emoji::CHECKMARK),
69            Err(_) => handle_output(&pb, self.error_message, "red", emoji::CROSS),
70        }
71
72        if self.clear {
73            pb.finish_and_clear()
74        } else {
75            pb.finish()
76        }
77
78        output
79    }
80
81    /// Spawn a progress bar with the given iterator and run the callback per element.
82    pub async fn progress_with<I, T, F, R, O>(&self, iter: I, cb: F) -> miette::Result<()>
83    where
84        I: IntoIterator<Item = T> + Clone,
85        F: Fn(T) -> R,
86        R: Future<Output = miette::Result<O>>,
87    {
88        let vec = iter.into_iter().collect::<Vec<_>>();
89        let len = vec.len();
90        let width = len.to_string().len();
91        let pb = ProgressBar::new(len as u64);
92        pb.set_style(
93            ProgressStyle::with_template(&format!(
94                "[{{pos:>{width}}}/{{len:{width}}}] {{bar:.blue/white}} {{msg}}"
95            ))
96            .unwrap(),
97        );
98
99        if let Some(message) = self.message {
100            pb.set_message(message.to_owned());
101        }
102
103        for item in vec {
104            cb(item).await?;
105            pb.inc(1);
106        }
107
108        handle_output(&pb, self.success_message, "green", emoji::CHECKMARK);
109
110        if self.clear {
111            pb.finish_and_clear()
112        } else {
113            pb.finish()
114        }
115
116        Ok(())
117    }
118}
119
120/// Show the output message with custom color and emoji prefix after progress has finished.
121fn handle_output(pb: &ProgressBar, msg: Option<&str>, color: &str, prefix: Emoji) {
122    pb.set_style(
123        ProgressStyle::with_template(&format!("{{prefix:.{color}}} {{msg:.{color}}}")).unwrap(),
124    );
125    pb.set_prefix(format!("{}", prefix));
126    if let Some(msg) = msg {
127        pb.set_message(msg.to_owned());
128    }
129}