novel_cli/utils/
progress.rs

1use std::env;
2use std::fmt::Write;
3use std::io::{self, IsTerminal, Stderr};
4use std::sync::{Arc, Mutex};
5
6use color_eyre::eyre::Result;
7use indicatif::{ProgressState, ProgressStyle};
8use osc94::Progress;
9
10#[must_use]
11#[derive(Clone)]
12pub struct ProgressBar {
13    pb: Option<indicatif::ProgressBar>,
14    pb_osc94: Option<Arc<Mutex<Progress<Stderr>>>>,
15    message_width: usize,
16}
17
18impl ProgressBar {
19    pub fn new(total_size: u64) -> Result<Self> {
20        let pb = if ProgressBar::is_terminal() {
21            let pb = indicatif::ProgressBar::new(total_size);
22            pb.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {pos:>7}/{len:7} {msg:40} ({eta})")
23              .unwrap()
24              .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()}).progress_chars("#>-"));
25
26            Some(pb)
27        } else {
28            None
29        };
30
31        let pb_osc94 = if ProgressBar::is_terminal() && ProgressBar::supports_term_integration() {
32            let mut progress = Progress::default();
33            progress.start();
34
35            Some(Arc::new(Mutex::new(progress)))
36        } else {
37            None
38        };
39
40        let message_width = (super::terminal_size().0 / 3) as usize;
41
42        Ok(Self {
43            pb,
44            pb_osc94,
45            message_width,
46        })
47    }
48
49    pub fn inc<T>(&mut self, msg: T, chunk_size: usize) -> Result<()>
50    where
51        T: AsRef<str>,
52    {
53        if let Some(pb) = &self.pb {
54            pb.set_message(
55                console::truncate_str(msg.as_ref(), self.message_width, "...").to_string(),
56            );
57            pb.inc(chunk_size as u64);
58        }
59
60        if let Some(pb) = &self.pb
61            && let Some(pb_osc94) = &self.pb_osc94
62        {
63            let progress = (pb.position() as f64 / pb.length().unwrap() as f64 * 100.0) as u8;
64            pb_osc94.lock().unwrap().progress(progress).flush()?;
65        }
66
67        Ok(())
68    }
69
70    pub fn finish(&self) -> Result<()> {
71        if let Some(pb) = &self.pb {
72            pb.finish_and_clear();
73        }
74
75        if let Some(pb_osc94) = &self.pb_osc94 {
76            pb_osc94.lock().unwrap().hidden().flush()?;
77        }
78
79        Ok(())
80    }
81
82    fn supports_term_integration() -> bool {
83        let windows_terminal = env::var("WT_SESSION").is_ok();
84        let conemu = env::var("ConEmuANSI").ok() == Some("ON".into());
85        let wezterm = env::var("TERM_PROGRAM").ok() == Some("WezTerm".into());
86        let ghostty = env::var("TERM_PROGRAM").ok() == Some("ghostty".into());
87
88        windows_terminal || conemu || wezterm || ghostty
89    }
90
91    fn is_terminal() -> bool {
92        io::stdout().is_terminal()
93    }
94}
novel_cli/utils/progress.rs

novel_cli/utils/
progress.rs
