Skip to main content

zzz_arc/
progress.rs

1//! progress reporting functionality
2
3use 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    /// create new byte-based progress reporter, only shows progress if enabled
48    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    /// create new item-count progress reporter, only shows progress if enabled
59    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    /// update progress with current bytes processed
70    pub fn update(&self, processed_bytes: u64) {
71        if let Some(bar) = &self.bar {
72            bar.set_position(processed_bytes);
73        }
74    }
75
76    /// set the current position
77    pub fn set_position(&self, pos: u64) {
78        if let Some(bar) = &self.bar {
79            bar.set_position(pos);
80        }
81    }
82
83    /// set the total length
84    pub fn set_length(&self, len: u64) {
85        if let Some(bar) = &self.bar {
86            bar.set_length(len);
87        }
88    }
89
90    /// finish progress with a message
91    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    /// finish progress and clear
98    pub fn finish(&self) {
99        if let Some(bar) = &self.bar {
100            bar.finish_and_clear();
101        }
102    }
103
104    /// check if verbose logging is enabled
105    pub fn is_verbose(&self) -> bool {
106        self.verbose
107    }
108
109    /// check if this progress tracks items (vs bytes)
110    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        // Should not panic when progress is disabled
153        progress.update(500);
154        progress.finish_with_message("test");
155        progress.finish();
156
157        // Progress bar should be None when disabled
158        assert!(progress.bar.is_none());
159    }
160
161    #[test]
162    fn test_progress_enabled() {
163        let progress = Progress::new(true, 1000, false);
164
165        // Progress bar should exist when enabled
166        assert!(progress.bar.is_some());
167
168        // Should not panic with progress operations
169        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        // Should handle zero total bytes without panic
179        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        // Should handle updates beyond total without panic
188        progress.update(150);
189        progress.finish();
190    }
191
192    #[test]
193    fn test_progress_is_verbose() {
194        // Enabled progress should be verbose
195        let progress_enabled = Progress::new(true, 1000, true);
196        assert!(progress_enabled.is_verbose());
197
198        // Disabled progress should not be verbose
199        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}