Skip to main content

rusty_rich/
progress_columns.rs

1//! Progress column types — equivalent to Python Rich's progress column
2//! system (SpinnerColumn, BarColumn, TextColumn, etc.).
3
4use std::time::Instant;
5
6use crate::progress::{ProgressBar, Task};
7use crate::spinner::Spinner;
8use crate::style::{Style, StyleType};
9
10// ---------------------------------------------------------------------------
11// ProgressColumn trait
12// ---------------------------------------------------------------------------
13
14/// A column in a progress display. Each column renders one cell per task.
15pub trait ProgressColumn: std::fmt::Debug {
16    /// Render this column for the given task into a string.
17    fn render(&self, task: &Task, width: usize, elapsed: std::time::Duration) -> String;
18}
19
20// ---------------------------------------------------------------------------
21// TextColumn
22// ---------------------------------------------------------------------------
23
24/// Displays a formatted text field. The text is taken from `task.fields["key"]`
25/// and formatted with the given format string.
26#[derive(Debug, Clone)]
27pub struct TextColumn {
28    /// Key into the task's `fields` HashMap.
29    pub key: String,
30    /// Format string (e.g. "{:>10}").
31    pub format: String,
32    /// Style for the text.
33    pub style: Style,
34}
35
36impl TextColumn {
37    pub fn new(key: impl Into<String>) -> Self {
38        Self { key: key.into(), format: "{:>11}".to_string(), style: Style::new() }
39    }
40
41    pub fn format(mut self, fmt: impl Into<String>) -> Self { self.format = fmt.into(); self }
42    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
43}
44
45impl ProgressColumn for TextColumn {
46    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
47        let value = task.fields.get(&self.key).map(|s| s.as_str()).unwrap_or("?");
48        // Simple: just return the value (formatting could use format args but
49        // we keep it simple)
50        let ansi = self.style.to_ansi();
51        let reset = self.style.reset_ansi();
52        format!("{ansi}{value}{reset}")
53    }
54}
55
56// ---------------------------------------------------------------------------
57// BarColumn
58// ---------------------------------------------------------------------------
59
60/// Renders the progress bar itself.
61#[derive(Debug, Clone)]
62pub struct BarColumn {
63    /// The underlying progress bar template.
64    pub bar: ProgressBar,
65    /// Width override (None = auto from available space).
66    pub width: Option<usize>,
67}
68
69impl BarColumn {
70    pub fn new() -> Self {
71        Self { bar: ProgressBar::new(), width: None }
72    }
73
74    pub fn complete_style(mut self, s: Style) -> Self { self.bar = self.bar.complete_style(s); self }
75    pub fn finished_style(mut self, s: Style) -> Self { self.bar = self.bar.remaining_style(s); self }
76    pub fn width(mut self, w: usize) -> Self { self.width = Some(w); self }
77}
78
79impl ProgressColumn for BarColumn {
80    fn render(&self, task: &Task, width: usize, _elapsed: std::time::Duration) -> String {
81        let w = self.width.unwrap_or(width.saturating_sub(2));
82        let mut bar = self.bar.clone();
83        bar.total = task.total;
84        bar.completed = task.completed;
85        bar.width = Some(w);
86        bar.render(w)
87    }
88}
89
90impl Default for BarColumn {
91    fn default() -> Self { Self::new() }
92}
93
94// ---------------------------------------------------------------------------
95// SpinnerColumn
96// ---------------------------------------------------------------------------
97
98/// Shows a spinner for tasks that are not finished, and "✓" when complete.
99#[derive(Debug, Clone)]
100pub struct SpinnerColumn {
101    pub spinner: Spinner,
102    pub style: Style,
103    pub finished_style: Style,
104    pub finished_text: String,
105}
106
107impl SpinnerColumn {
108    pub fn new() -> Self {
109        Self {
110            spinner: Spinner::default(),
111            style: Style::new(),
112            finished_style: Style::new().color(crate::color::Color::parse("green").unwrap()).bold(true),
113            finished_text: "✓".to_string(),
114        }
115    }
116
117    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
118    pub fn finished_style(mut self, s: Style) -> Self { self.finished_style = s; self }
119}
120
121impl ProgressColumn for SpinnerColumn {
122    fn render(&self, task: &Task, _width: usize, elapsed: std::time::Duration) -> String {
123        if task.is_finished() {
124            let a = self.finished_style.to_ansi();
125            let r = self.finished_style.reset_ansi();
126            format!("{a}{}{r}", self.finished_text)
127        } else {
128            let frame = self.spinner.frame_at(elapsed);
129            let a = self.style.to_ansi();
130            let r = self.style.reset_ansi();
131            format!("{a}{frame}{r}")
132        }
133    }
134}
135
136impl Default for SpinnerColumn {
137    fn default() -> Self { Self::new() }
138}
139
140// ---------------------------------------------------------------------------
141// TimeElapsedColumn
142// ---------------------------------------------------------------------------
143
144/// Shows elapsed time since task started.
145#[derive(Debug, Clone)]
146pub struct TimeElapsedColumn {
147    pub style: Style,
148    pub paused_style: Style,
149}
150
151impl TimeElapsedColumn {
152    pub fn new() -> Self {
153        Self { style: Style::new(), paused_style: Style::new().dim(true) }
154    }
155}
156
157impl ProgressColumn for TimeElapsedColumn {
158    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
159        let d = task.elapsed();
160        let s = format_duration_short(&d);
161        let a = self.style.to_ansi();
162        let r = self.style.reset_ansi();
163        format!("{a}{s}{r}")
164    }
165}
166
167impl Default for TimeElapsedColumn {
168    fn default() -> Self { Self::new() }
169}
170
171// ---------------------------------------------------------------------------
172// TimeRemainingColumn
173// ---------------------------------------------------------------------------
174
175/// Shows estimated time remaining.
176#[derive(Debug, Clone)]
177pub struct TimeRemainingColumn {
178    pub style: Style,
179    pub elapsed_when_finished: bool,
180}
181
182impl TimeRemainingColumn {
183    pub fn new() -> Self {
184        Self { style: Style::new(), elapsed_when_finished: false }
185    }
186}
187
188impl ProgressColumn for TimeRemainingColumn {
189    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
190        let text = if task.is_finished() {
191            if self.elapsed_when_finished {
192                format_duration_short(&task.elapsed())
193            } else {
194                String::new()
195            }
196        } else {
197            task.time_remaining()
198                .map(|d| format_duration_short(&d))
199                .unwrap_or_else(|| "?".to_string())
200        };
201
202        let a = self.style.to_ansi();
203        let r = self.style.reset_ansi();
204        format!("{a}{text}{r}")
205    }
206}
207
208impl Default for TimeRemainingColumn {
209    fn default() -> Self { Self::new() }
210}
211
212// ---------------------------------------------------------------------------
213// TaskProgressColumn
214// ---------------------------------------------------------------------------
215
216/// Shows percentage complete as text.
217#[derive(Debug, Clone)]
218pub struct TaskProgressColumn {
219    pub style: Style,
220}
221
222impl TaskProgressColumn {
223    pub fn new() -> Self { Self { style: Style::new() } }
224
225    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
226}
227
228impl ProgressColumn for TaskProgressColumn {
229    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
230        if task.total.is_some() {
231            let pct = (task.progress() * 100.0) as usize;
232            let s = format!("{pct:>3}%");
233            let a = self.style.to_ansi();
234            let r = self.style.reset_ansi();
235            format!("{a}{s}{r}")
236        } else {
237            String::new()
238        }
239    }
240}
241
242impl Default for TaskProgressColumn {
243    fn default() -> Self { Self::new() }
244}
245
246// ---------------------------------------------------------------------------
247// MofNCompleteColumn
248// ---------------------------------------------------------------------------
249
250/// Shows "completed / total" style.
251#[derive(Debug, Clone)]
252pub struct MofNCompleteColumn {
253    pub style: Style,
254    pub separator: String,
255}
256
257impl MofNCompleteColumn {
258    pub fn new() -> Self {
259        Self { style: Style::new(), separator: "/".to_string() }
260    }
261}
262
263impl ProgressColumn for MofNCompleteColumn {
264    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
265        let completed = task.completed as usize;
266        if let Some(total) = task.total {
267            let total = total as usize;
268            let s = format!("{completed}{}{total}", self.separator);
269            let a = self.style.to_ansi();
270            let r = self.style.reset_ansi();
271            format!("{a}{s}{r}")
272        } else {
273            format!("{completed}")
274        }
275    }
276}
277
278impl Default for MofNCompleteColumn {
279    fn default() -> Self { Self::new() }
280}
281
282// ---------------------------------------------------------------------------
283// Helper: short duration format
284// ---------------------------------------------------------------------------
285
286fn format_duration_short(d: &std::time::Duration) -> String {
287    let secs = d.as_secs();
288    if secs < 60 {
289        format!("0:{secs:02}")
290    } else if secs < 3600 {
291        format!("{}:{:02}", secs / 60, secs % 60)
292    } else {
293        format!("{}:{:02}:{:02}", secs / 3600, (secs % 3600) / 60, secs % 60)
294    }
295}
296
297// ---------------------------------------------------------------------------
298// Helper: file size formatting
299// ---------------------------------------------------------------------------
300
301/// Format bytes into human-readable form using decimal (1000-based) units.
302pub fn format_size(bytes: f64) -> String {
303    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB", "PB"];
304    let mut value = bytes;
305    let mut unit_idx = 0;
306    while value >= 1000.0 && unit_idx < UNITS.len() - 1 {
307        value /= 1000.0;
308        unit_idx += 1;
309    }
310    if unit_idx == 0 {
311        format!("{:.0} {}", value, UNITS[unit_idx])
312    } else {
313        format!("{:.1} {}", value, UNITS[unit_idx])
314    }
315}
316
317/// Format a transfer speed (bytes per second) into human-readable form.
318pub fn format_speed(bytes_per_sec: f64) -> String {
319    format!("{}/s", format_size(bytes_per_sec))
320}
321
322// ---------------------------------------------------------------------------
323// FileSizeColumn
324// ---------------------------------------------------------------------------
325
326/// Shows the completed file size in human-readable format.
327#[derive(Debug, Clone)]
328pub struct FileSizeColumn {
329    pub style: Style,
330}
331
332impl FileSizeColumn {
333    pub fn new() -> Self {
334        Self { style: Style::new() }
335    }
336
337    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
338}
339
340impl ProgressColumn for FileSizeColumn {
341    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
342        let size = format_size(task.completed);
343        let a = self.style.to_ansi();
344        let r = self.style.reset_ansi();
345        format!("{a}{size}{r}")
346    }
347}
348
349impl Default for FileSizeColumn {
350    fn default() -> Self { Self::new() }
351}
352
353// ---------------------------------------------------------------------------
354// TotalFileSizeColumn
355// ---------------------------------------------------------------------------
356
357/// Shows the total file size in human-readable format.
358#[derive(Debug, Clone)]
359pub struct TotalFileSizeColumn {
360    pub style: Style,
361}
362
363impl TotalFileSizeColumn {
364    pub fn new() -> Self {
365        Self { style: Style::new() }
366    }
367
368    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
369}
370
371impl ProgressColumn for TotalFileSizeColumn {
372    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
373        let a = self.style.to_ansi();
374        let r = self.style.reset_ansi();
375        if let Some(total) = task.total {
376            let size = format_size(total);
377            format!("{a}{size}{r}")
378        } else {
379            String::new()
380        }
381    }
382}
383
384impl Default for TotalFileSizeColumn {
385    fn default() -> Self { Self::new() }
386}
387
388// ---------------------------------------------------------------------------
389// DownloadColumn
390// ---------------------------------------------------------------------------
391
392/// Shows "completed/total" with file size formatting.
393#[derive(Debug, Clone)]
394pub struct DownloadColumn {
395    pub style: Style,
396    pub separator: String,
397}
398
399impl DownloadColumn {
400    pub fn new() -> Self {
401        Self { style: Style::new(), separator: "/".to_string() }
402    }
403
404    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
405    pub fn separator(mut self, sep: impl Into<String>) -> Self { self.separator = sep.into(); self }
406}
407
408impl ProgressColumn for DownloadColumn {
409    fn render(&self, task: &Task, _width: usize, _elapsed: std::time::Duration) -> String {
410        let a = self.style.to_ansi();
411        let r = self.style.reset_ansi();
412        let completed = format_size(task.completed);
413        if let Some(total) = task.total {
414            let total = format_size(total);
415            format!("{a}{completed}{}{total}{r}", self.separator)
416        } else {
417            format!("{a}{completed}{r}")
418        }
419    }
420}
421
422impl Default for DownloadColumn {
423    fn default() -> Self { Self::new() }
424}
425
426// ---------------------------------------------------------------------------
427// TransferSpeedColumn
428// ---------------------------------------------------------------------------
429
430/// Shows transfer speed in human-readable format (e.g., "1.5 MB/s").
431#[derive(Debug, Clone)]
432pub struct TransferSpeedColumn {
433    pub style: Style,
434}
435
436impl TransferSpeedColumn {
437    pub fn new() -> Self {
438        Self { style: Style::new() }
439    }
440
441    pub fn style(mut self, s: Style) -> Self { self.style = s; self }
442}
443
444impl ProgressColumn for TransferSpeedColumn {
445    fn render(&self, task: &Task, _width: usize, elapsed: std::time::Duration) -> String {
446        let secs = elapsed.as_secs_f64();
447        let a = self.style.to_ansi();
448        let r = self.style.reset_ansi();
449        if secs > 0.0 && task.completed > 0.0 {
450            let speed = task.completed / secs;
451            let s = format_speed(speed);
452            format!("{a}{s}{r}")
453        } else {
454            format!("{a}0 B/s{r}")
455        }
456    }
457}
458
459impl Default for TransferSpeedColumn {
460    fn default() -> Self { Self::new() }
461}
462
463// ---------------------------------------------------------------------------
464// Tests
465// ---------------------------------------------------------------------------
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470    use crate::progress::Task;
471
472    #[test]
473    fn test_text_column() {
474        let col = TextColumn::new("name");
475        let task = {
476            let mut t = Task::new(1, "test", Some(100.0));
477            t.fields.insert("name".into(), "Alice".into());
478            t
479        };
480        let result = col.render(&task, 20, std::time::Duration::from_secs(5));
481        assert!(result.contains("Alice"));
482    }
483
484    #[test]
485    fn test_spinner_column() {
486        let col = SpinnerColumn::new();
487        let task = Task::new(1, "test", Some(100.0));
488        let result = col.render(&task, 10, std::time::Duration::from_secs(1));
489        assert!(!result.is_empty());
490    }
491
492    #[test]
493    fn test_task_progress_column() {
494        let col = TaskProgressColumn::new();
495        let mut task = Task::new(1, "test", Some(100.0));
496        task.completed = 42.0;
497        let result = col.render(&task, 10, std::time::Duration::new(0, 0));
498        assert!(result.contains("42%"));
499    }
500
501    #[test]
502    fn test_format_size() {
503        assert_eq!(format_size(0.0), "0 B");
504        assert_eq!(format_size(500.0), "500 B");
505        assert_eq!(format_size(1500.0), "1.5 KB");
506        assert_eq!(format_size(2_500_000.0), "2.5 MB");
507    }
508
509    #[test]
510    fn test_format_speed() {
511        assert_eq!(format_speed(0.0), "0 B/s");
512        assert_eq!(format_speed(1500.0), "1.5 KB/s");
513    }
514
515    #[test]
516    fn test_file_size_column() {
517        let col = FileSizeColumn::new();
518        let mut task = Task::new(1, "test", Some(1000.0));
519        task.completed = 500.0;
520        let result = col.render(&task, 10, std::time::Duration::new(0, 0));
521        assert!(result.contains("500 B"));
522    }
523
524    #[test]
525    fn test_total_file_size_column() {
526        let col = TotalFileSizeColumn::new();
527        let task = Task::new(1, "test", Some(2_500_000.0));
528        let result = col.render(&task, 10, std::time::Duration::new(0, 0));
529        assert!(result.contains("2.5 MB"));
530    }
531
532    #[test]
533    fn test_download_column() {
534        let col = DownloadColumn::new();
535        let mut task = Task::new(1, "test", Some(1_500_000.0));
536        task.completed = 500_000.0;
537        let result = col.render(&task, 10, std::time::Duration::new(0, 0));
538        assert!(result.contains("500.0 KB"));
539        assert!(result.contains("1.5 MB"));
540    }
541
542    #[test]
543    fn test_transfer_speed_column() {
544        let col = TransferSpeedColumn::new();
545        let mut task = Task::new(1, "test", Some(1000.0));
546        task.completed = 500.0;
547        let result = col.render(&task, 10, std::time::Duration::from_secs(1));
548        assert!(result.contains("500 B/s"));
549    }
550}