1use indicatif::{ProgressBar, ProgressStyle};
4use std::io::Read;
5
6const PROGRESS_BYTES_TEMPLATE: &str =
7 "{spinner:.green} [{elapsed_precise}] [{bar:.cyan/blue}] {bytes}/{total_bytes} {bytes_per_sec} ({eta})";
8const PROGRESS_ITEMS_TEMPLATE: &str =
9 "{spinner:.green} [{elapsed_precise}] [{bar:.cyan/blue}] {pos}/{len} items ({eta})";
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum ProgressKind {
13 Bytes,
14 Items,
15}
16
17pub struct Progress {
18 bar: Option<ProgressBar>,
19 verbose: bool,
20 kind: ProgressKind,
21}
22
23impl Progress {
24 fn new_with_template(
25 enabled: bool,
26 total: u64,
27 template: &str,
28 verbose: bool,
29 kind: ProgressKind,
30 ) -> Self {
31 let bar = if enabled {
32 let pb = ProgressBar::new(total);
33 pb.set_style(
34 ProgressStyle::default_bar()
35 .template(template)
36 .expect("invalid progress template")
37 .progress_chars("█▉▊▋▌▍▎▏ "),
38 );
39 Some(pb)
40 } else {
41 None
42 };
43
44 Self { bar, verbose, kind }
45 }
46
47 pub fn new(enabled: bool, total_bytes: u64, verbose: bool) -> Self {
49 Self::new_with_template(
50 enabled,
51 total_bytes,
52 PROGRESS_BYTES_TEMPLATE,
53 verbose,
54 ProgressKind::Bytes,
55 )
56 }
57
58 pub fn new_items(enabled: bool, total_items: u64, verbose: bool) -> Self {
60 Self::new_with_template(
61 enabled,
62 total_items,
63 PROGRESS_ITEMS_TEMPLATE,
64 verbose,
65 ProgressKind::Items,
66 )
67 }
68
69 pub fn update(&self, processed_bytes: u64) {
71 if let Some(bar) = &self.bar {
72 bar.set_position(processed_bytes);
73 }
74 }
75
76 pub fn set_position(&self, pos: u64) {
78 if let Some(bar) = &self.bar {
79 bar.set_position(pos);
80 }
81 }
82
83 pub fn set_length(&self, len: u64) {
85 if let Some(bar) = &self.bar {
86 bar.set_length(len);
87 }
88 }
89
90 pub fn finish_with_message(&self, msg: &str) {
92 if let Some(bar) = &self.bar {
93 bar.finish_with_message(msg.to_string());
94 }
95 }
96
97 pub fn finish(&self) {
99 if let Some(bar) = &self.bar {
100 bar.finish_and_clear();
101 }
102 }
103
104 pub fn is_verbose(&self) -> bool {
106 self.verbose
107 }
108
109 pub fn is_items(&self) -> bool {
111 self.kind == ProgressKind::Items
112 }
113}
114
115pub struct ProgressReader<'a, R> {
116 inner: R,
117 progress: Option<&'a Progress>,
118 bytes_read: u64,
119}
120
121impl<'a, R> ProgressReader<'a, R> {
122 pub fn new(inner: R, progress: Option<&'a Progress>) -> Self {
123 Self {
124 inner,
125 progress,
126 bytes_read: 0,
127 }
128 }
129}
130
131impl<'a, R: Read> Read for ProgressReader<'a, R> {
132 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
133 let read = self.inner.read(buf)?;
134 if read > 0 {
135 self.bytes_read += read as u64;
136 if let Some(progress) = self.progress {
137 progress.update(self.bytes_read);
138 }
139 }
140 Ok(read)
141 }
142}
143
144#[cfg(test)]
145mod tests {
146 use super::*;
147
148 #[test]
149 fn test_progress_disabled() {
150 let progress = Progress::new(false, 1000, false);
151
152 progress.update(500);
154 progress.finish_with_message("test");
155 progress.finish();
156
157 assert!(progress.bar.is_none());
159 }
160
161 #[test]
162 fn test_progress_enabled() {
163 let progress = Progress::new(true, 1000, false);
164
165 assert!(progress.bar.is_some());
167
168 progress.update(500);
170 progress.update(750);
171 progress.finish_with_message("completed");
172 }
173
174 #[test]
175 fn test_progress_zero_total() {
176 let progress = Progress::new(true, 0, false);
177
178 progress.update(0);
180 progress.finish();
181 }
182
183 #[test]
184 fn test_progress_update_beyond_total() {
185 let progress = Progress::new(true, 100, false);
186
187 progress.update(150);
189 progress.finish();
190 }
191
192 #[test]
193 fn test_progress_is_verbose() {
194 let progress_enabled = Progress::new(true, 1000, true);
196 assert!(progress_enabled.is_verbose());
197
198 let progress_disabled = Progress::new(false, 1000, false);
200 assert!(!progress_disabled.is_verbose());
201 }
202
203 #[test]
204 fn test_progress_verbose_without_bar() {
205 let progress = Progress::new(false, 1000, true);
206 assert!(progress.is_verbose());
207 assert!(progress.bar.is_none());
208 }
209}