mobot/
progress.rs

1use crate::{api, Event};
2
3/// Represent the current state of the progressbar.
4enum ProgressState<'a> {
5    /// The task is still in progress.
6    Working,
7
8    // The task completed successfully -- show the str.
9    Done(&'a str),
10
11    /// The task failed -- show the str.
12    Failed(&'a str),
13}
14
15/// This method generates the progress bar string out of unicode block characters.
16fn progress_str(i: i64, state: ProgressState) -> String {
17    // Set of horizontal block characters of increasing size.
18    let blocks = vec![
19        '\u{258F}', '\u{258E}', '\u{258D}', '\u{258C}', '\u{258B}', '\u{258A}', '\u{2589}',
20    ];
21
22    let num_full_blocks = i / blocks.len() as i64;
23
24    let bar = blocks[blocks.len() - 1]
25        .to_string()
26        .repeat(num_full_blocks as usize);
27
28    let partial_bar = blocks[(i % blocks.len() as i64) as usize];
29
30    // If the state is Done or Failed, then show an additional message. Defaults to
31    // a unicode checkmark or cross.
32    match state {
33        ProgressState::Working => format!("{}{}", bar, partial_bar),
34        ProgressState::Done(c) => format!("{}{} {}", bar, partial_bar, c),
35        ProgressState::Failed(c) => format!("{}{} {}", bar, partial_bar, c),
36    }
37}
38
39/// A ProgressBar that can be rendered in a telegram message. This calls a long running async
40/// task and shows a progress bar while the task is running. The progress bar is updated every
41/// `update_interval` seconds. If the task completes before the `timeout` then the progress bar
42/// is replaced with a checkmark. If the task fails, then the progress bar is replaced with a
43/// cross.
44#[derive(Debug, Clone)]
45pub struct ProgressBar {
46    /// The timeout for the progress bar. If the task completes before this timeout, then the
47    /// progress bar is replaced with a checkmark.
48    pub timeout: std::time::Duration,
49
50    /// The update interval for the progress bar. The progress bar is updated every
51    /// `update_interval` seconds.
52    pub update_interval: std::time::Duration,
53
54    /// The string to show when the task fails.
55    pub failed_str: String,
56
57    /// The string to show when the task completes successfully.
58    pub done_str: String,
59
60    /// If true, then show the result of the task after the progress bar.
61    pub show_result: bool,
62}
63
64impl Default for ProgressBar {
65    fn default() -> Self {
66        Self {
67            timeout: std::time::Duration::from_secs(60),
68            update_interval: std::time::Duration::from_millis(500),
69            failed_str: '\u{2718}'.into(),
70            done_str: '\u{2714}'.into(),
71            show_result: false,
72        }
73    }
74}
75
76impl ProgressBar {
77    /// Create a new progress bar.
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Set the timeout for the progress bar. If the task completes before this timeout, then the
83    /// progress bar is replaced with a checkmark.
84    pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
85        self.timeout = timeout;
86        self
87    }
88
89    /// Set the update interval for the progress bar. The progress bar is updated every
90    /// `update_interval` seconds.
91    pub fn with_update_interval(mut self, update_interval: std::time::Duration) -> Self {
92        self.update_interval = update_interval;
93        self
94    }
95
96    /// Start the progress bar. This calls the async function `f` and shows a progress bar while
97    /// the task is running. The progress bar is updated every `update_interval` seconds. If the
98    /// task completes before the `timeout` then the progress bar is replaced with a checkmark.
99    /// If the task fails, then the progress bar is replaced with a cross.
100    ///
101    /// If `show_result` is true, then the result of the task is shown after the progress bar.
102    ///
103    /// Returns the result of the task.
104    pub async fn start<F, R>(&self, e: &Event, f: F) -> anyhow::Result<R>
105    where
106        F: futures::Future<Output = anyhow::Result<R>> + Send + 'static,
107        R: Default + Send + Sync + 'static,
108    {
109        // Send an empty message to get a message id for the progress bar.
110        let mut message = e.send_message("...").await?;
111
112        let (completed_tx, mut completed_rx) = tokio::sync::oneshot::channel();
113        tokio::spawn(async {
114            _ = completed_tx.send(f.await);
115        });
116
117        let mut count = 0;
118        let mut done = false;
119        let mut result: R = R::default();
120
121        while !done {
122            tokio::select! {
123                // Update the progress bar.
124                _ = tokio::time::sleep(self.update_interval) => {
125                    count += 1;
126                    message = e.edit_message(message.message_id, progress_str(count, ProgressState::Working)).await?;
127                    e.send_chat_action(api::ChatAction::Typing).await?;
128                }
129
130                // Timeout.
131                _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {
132                    done = true;
133                    message = e.edit_message(message.message_id,
134                        format!("{} {}", progress_str(count, ProgressState::Failed(self.failed_str.as_str())),
135                            "Something's wrong!")).await?;
136                }
137
138                // The future has completed.
139                v = &mut completed_rx => {
140                    done = true;
141                    result = v??;
142                    message = e.edit_message(message.message_id,
143                        progress_str(count, ProgressState::Done(self.done_str.as_str()))).await?;
144                    e.delete_message(message.message_id).await?;
145
146                }
147            }
148        }
149
150        Ok(result)
151    }
152}