Skip to main content

lunar_lib/progress/
cli_progress_bar.rs

1#[cfg(feature = "smarter_progress_bar")]
2use std::fmt::Arguments;
3use std::{
4    io::{self, Write},
5    sync::RwLock,
6};
7
8use crate::progress::ProgressRenderer;
9
10fn terminal_width() -> usize {
11    #[cfg(feature = "smarter_progress_bar")]
12    {
13        crossterm::terminal::size().map(|(w, _)| w).unwrap_or(80) as usize
14    }
15
16    #[cfg(not(feature = "smarter_progress_bar"))]
17    {
18        80
19    }
20}
21
22/// A default implementation of a [`ProgressRenderer`] meant for CLI applications
23pub struct CliProgressRenderer {
24    bar_width: usize,
25    suffix: RwLock<String>,
26}
27
28impl CliProgressRenderer {
29    pub fn new(bar_width: usize) -> Self {
30        Self {
31            bar_width,
32            suffix: RwLock::new(String::new()),
33        }
34    }
35
36    fn draw(&self, value: usize, max: usize, suffix: &str) {
37        let max = max.max(1);
38        let value = value.min(max);
39
40        let value_width = max.to_string().len();
41        let prefix = format!("[ {value:>value_width$} / {max} ]");
42
43        let bar_width = self.bar_width.saturating_sub(prefix.len() + 3);
44
45        let progress = (value as f64 / max as f64 * bar_width as f64).round() as usize;
46
47        let has_tip = progress > 0 && value < max;
48        let body = progress.saturating_sub(has_tip as usize);
49
50        let truncated_suffix = {
51            let suffix_max_width = terminal_width().saturating_sub(self.bar_width + 1);
52            &suffix[..suffix.len().min(suffix_max_width)]
53        };
54
55        print!(
56            "\r{prefix} [{body_str}{tip_str}{none_str}] {truncated_suffix}",
57            body_str = "=".repeat(body),
58            tip_str = if has_tip { ">" } else { "" },
59            none_str = "-".repeat(bar_width.saturating_sub(progress)) // Also print suffix, making sure it can never go past the `terminal_width()`
60        );
61        let _ = io::stdout().flush();
62    }
63
64    /// Prints text above the progress bar. Use this instead of [`println!()`] when a bar is running
65    #[cfg(feature = "smarter_progress_bar")]
66    pub fn print_above(&self, message: Arguments) {
67        use crossterm::{cursor::MoveToColumn, execute, terminal::Clear};
68
69        let mut stdout = io::stdout();
70
71        let _ = execute!(
72            stdout,
73            Clear(crossterm::terminal::ClearType::CurrentLine),
74            MoveToColumn(0)
75        );
76
77        let _ = stdout.write_fmt(message);
78        let _ = stdout.write("\n".as_bytes());
79    }
80
81    /// Sets the suffix of the progress bar for when it next updates
82    pub fn set_suffix(&self, suffix: impl AsRef<str>) {
83        let mut m_suffix = self.suffix.write().unwrap();
84        *m_suffix = suffix.as_ref().to_string()
85    }
86}
87
88impl ProgressRenderer for CliProgressRenderer {
89    fn on_start(&self, value: usize, max: usize) {
90        let suffix = self.suffix.read().unwrap();
91        self.draw(value, max, &suffix);
92    }
93
94    fn on_update(&self, value: usize, max: usize) {
95        let suffix = self.suffix.read().unwrap();
96        self.draw(value, max, &suffix);
97    }
98
99    fn on_finish(&self) {
100        println!()
101    }
102
103    fn on_notify(&self, msg: String) {
104        #[cfg(feature = "smarter_progress_bar")]
105        self.print_above(format_args!("{msg}"));
106    }
107
108    fn set_label(&self, msg: &str) {
109        self.set_suffix(msg);
110    }
111}