novel_cli/utils/
progress.rs

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

novel_cli/utils/
progress.rs
