xacli_components/components/
progress.rs

1use std::{
2    io::{stdout, Write},
3    sync::mpsc::{self, Receiver, RecvTimeoutError, Sender},
4    time::Duration,
5};
6
7use crossterm::{
8    cursor, queue,
9    style::Print,
10    terminal::{self, ClearType},
11};
12
13use crate::Result;
14
15pub struct ProgressBar {
16    total: u64,
17    message: String,
18    rx: Receiver<ProgressEvent>,
19}
20
21pub struct ProgressHandle {
22    tx: Sender<ProgressEvent>,
23}
24
25pub enum ProgressEvent {
26    Increment(u64),
27    SetProgress(u64),
28    SetMessage(String),
29    Finish,
30}
31
32impl ProgressBar {
33    pub fn new(total: u64, message: impl Into<String>) -> (Self, ProgressHandle) {
34        let (tx, rx) = mpsc::channel();
35        let bar = Self {
36            total,
37            message: message.into(),
38            rx,
39        };
40        let handle = ProgressHandle { tx };
41        (bar, handle)
42    }
43
44    pub fn run(mut self) -> Result<()> {
45        let mut stdout = stdout();
46        let mut current = 0u64;
47
48        self.render(&mut stdout, current)?;
49
50        loop {
51            match self.rx.recv_timeout(Duration::from_millis(100)) {
52                Ok(ProgressEvent::Increment(n)) => {
53                    current = (current + n).min(self.total);
54                    self.render(&mut stdout, current)?;
55                }
56                Ok(ProgressEvent::SetProgress(n)) => {
57                    current = n.min(self.total);
58                    self.render(&mut stdout, current)?;
59                }
60                Ok(ProgressEvent::SetMessage(msg)) => {
61                    self.message = msg;
62                    self.render(&mut stdout, current)?;
63                }
64                Ok(ProgressEvent::Finish) => {
65                    current = self.total;
66                    self.render(&mut stdout, current)?;
67                    println!(); // 换行
68                    break;
69                }
70                Err(RecvTimeoutError::Timeout) => {
71                    // 定期刷新
72                    self.render(&mut stdout, current)?;
73                }
74                Err(RecvTimeoutError::Disconnected) => {
75                    println!(); // 换行
76                    break;
77                }
78            }
79        }
80
81        Ok(())
82    }
83
84    fn render(&self, stdout: &mut impl Write, current: u64) -> Result<()> {
85        let percent = if self.total > 0 {
86            (current as f64 / self.total as f64) * 100.0
87        } else {
88            0.0
89        };
90
91        let bar_width = 40;
92        let filled = if self.total > 0 {
93            ((current as f64 / self.total as f64) * bar_width as f64) as usize
94        } else {
95            0
96        };
97        let bar = format!("[{}{}]", "=".repeat(filled), " ".repeat(bar_width - filled));
98
99        queue!(
100            stdout,
101            cursor::MoveToColumn(0),
102            terminal::Clear(ClearType::CurrentLine),
103            Print(&self.message),
104            Print(" "),
105            Print(&bar),
106            Print(format!(" {}/{} ({:.1}%)", current, self.total, percent)),
107        )?;
108
109        stdout.flush()?;
110        Ok(())
111    }
112}
113
114impl ProgressHandle {
115    pub fn inc(&self, n: u64) {
116        let _ = self.tx.send(ProgressEvent::Increment(n));
117    }
118
119    pub fn set(&self, n: u64) {
120        let _ = self.tx.send(ProgressEvent::SetProgress(n));
121    }
122
123    pub fn set_message(&self, msg: impl Into<String>) {
124        let _ = self.tx.send(ProgressEvent::SetMessage(msg.into()));
125    }
126
127    pub fn finish(&self) {
128        let _ = self.tx.send(ProgressEvent::Finish);
129    }
130}