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