1use std::future::Future;
4
5use console::Emoji;
6use indicatif::{ProgressBar, ProgressStyle};
7use tokio::time::Duration;
8
9use crate::constants::emoji;
10
11#[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 pub fn new() -> Self {
23 Self::default()
24 }
25
26 pub fn message(&mut self, message: &'a str) -> &mut Self {
28 self.message = Some(message);
29 self
30 }
31
32 pub fn success_message(&mut self, message: &'a str) -> &mut Self {
34 self.success_message = Some(message);
35 self
36 }
37
38 pub fn error_message(&mut self, message: &'a str) -> &mut Self {
40 self.error_message = Some(message);
41 self
42 }
43
44 #[allow(dead_code)]
46 pub fn clear(&mut self, clear: bool) -> &mut Self {
47 self.clear = clear;
48 self
49 }
50
51 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 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
120fn 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}