Skip to main content

rich/
progress.rs

1//! Progress displays.
2//!
3//! Port of `rich/progress.py`'s display and task model: a grid of tasks, one
4//! row each, whose cells come from a list of [`ProgressColumn`]s, laid out as
5//! upstream's `make_tasks_table` does: a [`Table::grid`] with `padding=(0, 1)`,
6//! each column's table-column options, and `expand`.
7//!
8//! Time is read from an injectable clock ([`Progress::clock`], upstream's
9//! `get_time`), so elapsed time, speed, ETA and spinner frames are
10//! deterministic under test. The default clock is monotonic.
11//!
12//! [`Progress::start`] runs the display live on the [`Live`](crate::live::Live)
13//! refresh thread ([`LiveProgress`]); [`LiveProgress::track`] and [`track`]
14//! port `track()`. `TextColumn` format strings use [`pyformat`].
15//!
16//! `transient` and `disable` apply to the live display; [`LiveProgress::wrap_read`]
17//! and [`LiveProgress::open`] port `wrap_file` and `open`.
18
19use std::cell::RefCell;
20use std::collections::{BTreeMap, HashMap, VecDeque};
21use std::sync::Arc;
22
23use crate::console::{Console, ConsoleOptions, Justify};
24use crate::filesize;
25use crate::progress_bar::ProgressBar;
26use crate::protocol::Renderable;
27use crate::pyformat::{self, FormatValue};
28use crate::segment::Segment;
29use crate::spinner::Spinner;
30use crate::style::{Style, StyleType};
31use crate::table::{Cell, ColumnOptions, Table};
32use crate::text::Text;
33
34/// Upstream keeps at most this many speed samples per task (`deque(maxlen=1000)`).
35const MAX_SAMPLES: usize = 1000;
36
37/// A source of the current time in seconds. Upstream's `GetTimeCallable`.
38pub use crate::console::GetTime;
39
40use crate::console::monotonic;
41
42/// Identifies a task within one [`Progress`]. Upstream's `TaskID`.
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
44pub struct TaskId(pub usize);
45
46/// Estimated time remaining. Port of `TimeRemainingColumn`, including its
47/// half-second render cache (`max_refresh = 0.5`).
48pub struct TimeRemainingColumn {
49    compact: bool,
50    elapsed_when_finished: bool,
51    cache: RefCell<HashMap<TaskId, (f64, Text)>>,
52}
53
54impl TimeRemainingColumn {
55    /// `compact` drops the hours when there are none (`05:03`);
56    /// `elapsed_when_finished` shows the elapsed time once a task finishes.
57    pub fn new(compact: bool, elapsed_when_finished: bool) -> Self {
58        TimeRemainingColumn {
59            compact,
60            elapsed_when_finished,
61            cache: RefCell::new(HashMap::new()),
62        }
63    }
64}
65
66/// An animated spinner. Port of `SpinnerColumn`: one spinner shared by every
67/// row, whose animation starts at its first render.
68pub struct SpinnerColumn {
69    spinner: Spinner,
70    style: StyleType,
71    finished_text: String,
72}
73
74impl SpinnerColumn {
75    /// A spinner by name (e.g. `"dots"`), styled `progress.spinner`, showing
76    /// `finished_text` (console markup) once a task finishes.
77    pub fn new(name: &str, finished_text: impl Into<String>) -> Self {
78        SpinnerColumn {
79            spinner: Spinner::new(name),
80            style: StyleType::Name("progress.spinner".to_string()),
81            finished_text: finished_text.into(),
82        }
83    }
84
85    /// Animation speed multiplier (default 1.0).
86    pub fn speed(mut self, speed: f64) -> Self {
87        self.spinner = self.spinner.speed(speed);
88        self
89    }
90
91    /// Style of the spinner frame (default `progress.spinner`).
92    pub fn style(mut self, style: impl Into<StyleType>) -> Self {
93        self.style = style.into();
94        self
95    }
96}
97
98/// A text cell built from a format string. Port of `TextColumn`: the format
99/// is expanded against the task as `text_format.format(task=task)`, so it can
100/// use any task attribute (`{task.completed}`, `{task.percentage:>3.0f}`) and
101/// per-task fields (`{task.fields[name]}`).
102pub struct TextColumn {
103    text_format: String,
104    style: StyleType,
105    justify: Justify,
106    markup: bool,
107}
108
109impl TextColumn {
110    /// `TextColumn(text_format)` with upstream's defaults: no style, left
111    /// justified, console markup on.
112    pub fn new(text_format: impl Into<String>) -> Self {
113        TextColumn {
114            text_format: text_format.into(),
115            style: StyleType::default(),
116            justify: Justify::Left,
117            markup: true,
118        }
119    }
120
121    /// The style of the whole cell (upstream `style`).
122    pub fn style(mut self, style: impl Into<StyleType>) -> Self {
123        self.style = style.into();
124        self
125    }
126
127    /// Justify the text within the column (upstream `justify`).
128    pub fn justify(mut self, justify: Justify) -> Self {
129        self.justify = justify;
130        self
131    }
132
133    /// Parse the expanded text as console markup (upstream `markup`, default on).
134    pub fn markup(mut self, markup: bool) -> Self {
135        self.markup = markup;
136        self
137    }
138
139    fn render(&self, task: &Task) -> Text {
140        let expanded = pyformat::format(&self.text_format, |name| task.format_field(name));
141        let mut text = if self.markup {
142            Text::from_markup(&expanded).unwrap_or_else(|_| Text::new(expanded.clone()))
143        } else {
144            Text::new(expanded)
145        };
146        text.set_base_style(self.style.clone());
147        text.set_justify(self.justify);
148        text
149    }
150}
151
152/// A column in a [`Progress`] display. Mirrors upstream's `ProgressColumn`s.
153pub enum ProgressColumn {
154    /// The task description as console markup
155    /// (`TextColumn("[progress.description]{task.description}")`).
156    Description,
157    /// A static text cell with an explicit style (a simplified `TextColumn`).
158    Text(String, Style),
159    /// A text cell formatted from the task (`TextColumn`).
160    TextFormat(TextColumn),
161    /// The same renderable in every row (`RenderableColumn`). It may span
162    /// several lines; the row grows to fit it.
163    Renderable(Arc<dyn Renderable + Send + Sync>),
164    /// The flexing progress bar (`BarColumn`).
165    Bar,
166    /// The completion percentage `"{pct:>3}%"` (default `TaskProgressColumn`).
167    Percentage,
168    /// `TaskProgressColumn(show_speed=…)`: the percentage, or for a task with no
169    /// total and `show_speed`, the rate in `it/s`.
170    TaskProgress { show_speed: bool },
171    /// `"{completed}/{total}"` (`MofNCompleteColumn`, `progress.download`).
172    MofN,
173    /// `"{completed}/{total} {unit}"` in shared SI byte units, e.g. `0.5/1.0 kB`
174    /// (`DownloadColumn`, `progress.download`).
175    Download,
176    /// As [`Download`](Self::Download) in binary units (`DownloadColumn(binary_units=True)`).
177    BinaryDownload,
178    /// Elapsed time `H:MM:SS` (`TimeElapsedColumn`, `progress.elapsed`).
179    TimeElapsed,
180    /// Estimated time remaining (`TimeRemainingColumn`, `progress.remaining`).
181    TimeRemaining(TimeRemainingColumn),
182    /// Data speed, e.g. `1.2 MB/s` (`TransferSpeedColumn`, `progress.data.speed`).
183    TransferSpeed,
184    /// Completed size in decimal units (`FileSizeColumn`, `progress.filesize`).
185    FileSize,
186    /// Total size in decimal units (`TotalFileSizeColumn`, `progress.filesize.total`).
187    TotalFileSize,
188    /// An animated spinner (`SpinnerColumn`).
189    Spinner(SpinnerColumn),
190    /// A bar with its width and styles set (`BarColumn(bar_width=…, style=…)`).
191    /// A `bar_width` of `None` lets the bar fill its column.
192    BarWith(BarColumn),
193    /// A column with explicit table-column options (upstream's `table_column=`
194    /// argument). See [`ProgressColumn::with_table_column`].
195    WithTableColumn(Box<ProgressColumn>, ColumnOptions),
196    /// A user-defined column: upstream's subclass of `ProgressColumn`. See
197    /// [`CustomProgressColumn`] and [`ProgressColumn::custom`].
198    Custom(CustomColumn),
199}
200
201/// A user-defined progress column. Port of subclassing upstream's abstract
202/// `ProgressColumn`: implement [`render`](Self::render), and optionally the
203/// table column and `max_refresh`.
204pub trait CustomProgressColumn: Send + Sync {
205    /// The cell for `task`. Port of `ProgressColumn.render`.
206    fn render(&self, task: &Task) -> Cell;
207
208    /// The grid column this column occupies. Port of `get_table_column`
209    /// (upstream's default `Column()`).
210    fn table_column(&self) -> ColumnOptions {
211        ColumnOptions::default()
212    }
213
214    /// Reuse a render younger than this many seconds (upstream
215    /// `max_refresh`, default `None`: render every time). As upstream, the
216    /// cache is only consulted while the task has completed nothing.
217    fn max_refresh(&self) -> Option<f64> {
218        None
219    }
220}
221
222/// A [`CustomProgressColumn`] with upstream's `ProgressColumn.__call__`
223/// render cache.
224pub struct CustomColumn {
225    column: Arc<dyn CustomProgressColumn>,
226    cache: RefCell<HashMap<TaskId, (f64, Cell)>>,
227}
228
229impl CustomColumn {
230    /// Wrap a user-defined column.
231    pub fn new(column: Arc<dyn CustomProgressColumn>) -> Self {
232        CustomColumn {
233            column,
234            cache: RefCell::new(HashMap::new()),
235        }
236    }
237
238    /// The wrapped column.
239    pub fn column(&self) -> &Arc<dyn CustomProgressColumn> {
240        &self.column
241    }
242
243    /// Port of `ProgressColumn.__call__`.
244    fn call(&self, task: &Task) -> Cell {
245        let now = task.now();
246        let max_refresh = self.column.max_refresh();
247        if let Some(max_refresh) = max_refresh {
248            if task.completed == 0.0 {
249                if let Some((timestamp, cell)) = self.cache.borrow().get(&task.id) {
250                    if timestamp + max_refresh > now {
251                        return cell.clone();
252                    }
253                }
254            }
255        }
256        let cell = self.column.render(task);
257        self.cache.borrow_mut().insert(task.id, (now, cell.clone()));
258        cell
259    }
260}
261
262/// A progress bar column's width and styles. Port of `BarColumn`'s arguments.
263#[derive(Clone, Debug)]
264pub struct BarColumn {
265    bar_width: Option<usize>,
266    style: StyleType,
267    complete_style: StyleType,
268    finished_style: StyleType,
269    pulse_style: StyleType,
270}
271
272impl Default for BarColumn {
273    fn default() -> Self {
274        BarColumn {
275            bar_width: Some(40),
276            style: "bar.back".into(),
277            complete_style: "bar.complete".into(),
278            finished_style: "bar.finished".into(),
279            pulse_style: "bar.pulse".into(),
280        }
281    }
282}
283
284impl BarColumn {
285    /// Upstream's defaults: 40 cells wide, `bar.*` styles.
286    pub fn new() -> Self {
287        BarColumn::default()
288    }
289
290    /// The bar width, or `None` to fill the column (upstream `bar_width`).
291    pub fn bar_width(mut self, width: Option<usize>) -> Self {
292        self.bar_width = width;
293        self
294    }
295
296    /// The background style (upstream `style`).
297    pub fn style(mut self, style: impl Into<StyleType>) -> Self {
298        self.style = style.into();
299        self
300    }
301
302    /// The completed-part style (upstream `complete_style`).
303    pub fn complete_style(mut self, style: impl Into<StyleType>) -> Self {
304        self.complete_style = style.into();
305        self
306    }
307
308    /// The finished style (upstream `finished_style`).
309    pub fn finished_style(mut self, style: impl Into<StyleType>) -> Self {
310        self.finished_style = style.into();
311        self
312    }
313
314    /// The pulse style (upstream `pulse_style`).
315    pub fn pulse_style(mut self, style: impl Into<StyleType>) -> Self {
316        self.pulse_style = style.into();
317        self
318    }
319
320    /// Port of `BarColumn.render`.
321    fn render(&self, task: &Task) -> ProgressBar {
322        let bar = match task.total {
323            Some(total) => ProgressBar::new(total.max(0.0), task.completed.max(0.0)),
324            None => ProgressBar::indeterminate(),
325        };
326        let bar = match self.bar_width {
327            Some(width) => bar.width(width.max(1)),
328            None => bar,
329        };
330        bar.pulse(!task.started())
331            .animation_time(task.now())
332            .style(self.style.clone())
333            .complete_style(self.complete_style.clone())
334            .finished_style(self.finished_style.clone())
335            .pulse_style(self.pulse_style.clone())
336    }
337}
338
339impl ProgressColumn {
340    /// `TimeRemainingColumn()` with upstream's defaults.
341    pub fn time_remaining() -> Self {
342        ProgressColumn::TimeRemaining(TimeRemainingColumn::new(false, false))
343    }
344
345    /// `SpinnerColumn()` with upstream's defaults (`dots`, finished text `" "`).
346    pub fn spinner() -> Self {
347        ProgressColumn::Spinner(SpinnerColumn::new("dots", " "))
348    }
349
350    /// A user-defined column (upstream's `ProgressColumn` subclass).
351    pub fn custom(column: impl CustomProgressColumn + 'static) -> Self {
352        ProgressColumn::Custom(CustomColumn::new(Arc::new(column)))
353    }
354
355    /// This column with explicit table-column options, as upstream's
356    /// `table_column=Column(...)` argument sets them: width, ratio, justify,
357    /// wrapping and style of the grid column.
358    pub fn with_table_column(self, options: ColumnOptions) -> Self {
359        let inner = match self {
360            ProgressColumn::WithTableColumn(inner, _) => *inner,
361            column => column,
362        };
363        ProgressColumn::WithTableColumn(Box::new(inner), options)
364    }
365
366    /// Port of `get_table_column()`: text columns default to
367    /// `Column(no_wrap=True)`, the rest to `Column()`.
368    pub fn table_column(&self) -> ColumnOptions {
369        match self {
370            ProgressColumn::WithTableColumn(_, options) => options.clone(),
371            ProgressColumn::Custom(column) => column.column.table_column(),
372            ProgressColumn::Description
373            | ProgressColumn::Text(..)
374            | ProgressColumn::TextFormat(_)
375            | ProgressColumn::Percentage
376            | ProgressColumn::TaskProgress { .. } => ColumnOptions {
377                no_wrap: true,
378                ..ColumnOptions::default()
379            },
380            _ => ColumnOptions::default(),
381        }
382    }
383
384    /// The grid cell for `task`: the column's `__call__(task)`. Stateful
385    /// columns (spinners, the remaining-time and custom caches) advance as
386    /// upstream's do, so call it once per task per refresh.
387    pub fn table_cell(&self, task: &Task) -> Cell {
388        match self {
389            ProgressColumn::WithTableColumn(inner, _) => inner.table_cell(task),
390            ProgressColumn::Custom(column) => column.call(task),
391            ProgressColumn::Bar => Cell::Renderable(Arc::new(BarColumn::default().render(task))),
392            ProgressColumn::BarWith(column) => Cell::Renderable(Arc::new(column.render(task))),
393            ProgressColumn::Renderable(renderable) => Cell::Renderable(renderable.clone()),
394            column => Cell::Text(column.cell(task)),
395        }
396    }
397
398    /// The cell for `task` (never called on [`ProgressColumn::Bar`]).
399    fn cell(&self, task: &Task) -> Text {
400        let named = |plain: String, style: &str| Text::styled(plain, style);
401        match self {
402            // `TextColumn`s hand their `justify` (default left) to the text, so
403            // it overrides the table column's.
404            ProgressColumn::Description => {
405                let markup = format!("[progress.description]{}", task.description);
406                Text::from_markup(&markup)
407                    .unwrap_or_else(|_| Text::new(task.description.clone()))
408                    .justify(Justify::Left)
409            }
410            ProgressColumn::Text(text, style) => {
411                Text::styled(text.clone(), style.clone()).justify(Justify::Left)
412            }
413            ProgressColumn::TextFormat(column) => column.render(task),
414            ProgressColumn::Bar
415            | ProgressColumn::BarWith(_)
416            | ProgressColumn::Renderable(_)
417            | ProgressColumn::Custom(_)
418            | ProgressColumn::WithTableColumn(..) => {
419                unreachable!("bar, renderable and wrapped columns have no text cell")
420            }
421            ProgressColumn::Percentage => task.percentage_cell().justify(Justify::Left),
422            ProgressColumn::TaskProgress { show_speed } => {
423                if task.total.is_none() && *show_speed {
424                    render_speed(
425                        task.finished_speed
426                            .filter(|s| *s != 0.0)
427                            .or_else(|| task.speed()),
428                    )
429                } else {
430                    task.percentage_cell().justify(Justify::Left)
431                }
432            }
433            ProgressColumn::MofN => named(task.mofn_text(), "progress.download"),
434            ProgressColumn::Download => named(task.download_text(false), "progress.download"),
435            ProgressColumn::BinaryDownload => named(task.download_text(true), "progress.download"),
436            ProgressColumn::TimeElapsed => {
437                let elapsed = if task.finished() {
438                    task.finished_time
439                } else {
440                    task.elapsed()
441                };
442                let text = match elapsed {
443                    None => "-:--:--".to_string(),
444                    Some(elapsed) => timedelta(elapsed.max(0.0) as i64),
445                };
446                named(text, "progress.elapsed")
447            }
448            ProgressColumn::TimeRemaining(column) => column.render(task),
449            ProgressColumn::TransferSpeed => {
450                let speed = task
451                    .finished_speed
452                    .filter(|s| *s != 0.0)
453                    .or_else(|| task.speed());
454                let text = match speed {
455                    None => "?".to_string(),
456                    Some(speed) => format!("{}/s", filesize::decimal_signed(speed as i64)),
457                };
458                named(text, "progress.data.speed")
459            }
460            // `filesize.decimal(int(task.completed))`: `int()` truncates
461            // toward zero and keeps the sign.
462            ProgressColumn::FileSize => named(
463                filesize::decimal_signed(task.completed as i64),
464                "progress.filesize",
465            ),
466            ProgressColumn::TotalFileSize => named(
467                task.total
468                    .map_or_else(String::new, |total| filesize::decimal_signed(total as i64)),
469                "progress.filesize.total",
470            ),
471            ProgressColumn::Spinner(column) => {
472                if task.finished() {
473                    Text::from_markup(&column.finished_text)
474                        .unwrap_or_else(|_| Text::new(column.finished_text.clone()))
475                } else {
476                    // Upstream's `self.spinner.render(task.get_time())`: the
477                    // spinner itself starts its animation at the first render.
478                    let mut frame = column.spinner.render(task.now());
479                    frame.set_base_style(column.style.clone());
480                    frame
481                }
482            }
483        }
484    }
485}
486
487impl TimeRemainingColumn {
488    fn render(&self, task: &Task) -> Text {
489        // `ProgressColumn.__call__`: reuse a render younger than max_refresh,
490        // but only while the task has completed nothing (`not task.completed`).
491        let now = task.now();
492        if task.completed == 0.0 {
493            if let Some((timestamp, text)) = self.cache.borrow().get(&task.id) {
494                if timestamp + 0.5 > now {
495                    return text.clone();
496                }
497            }
498        }
499        let (task_time, style) = if self.elapsed_when_finished && task.finished() {
500            (task.finished_time, "progress.elapsed")
501        } else {
502            (task.time_remaining(), "progress.remaining")
503        };
504        let text = if task.total.is_none() {
505            Text::styled("", style)
506        } else {
507            match task_time {
508                None => Text::styled(if self.compact { "--:--" } else { "-:--:--" }, style),
509                Some(task_time) => {
510                    let whole = task_time as i64;
511                    let (minutes, seconds) = (whole.div_euclid(60), whole.rem_euclid(60));
512                    let (hours, minutes) = (minutes.div_euclid(60), minutes.rem_euclid(60));
513                    let formatted = if self.compact && hours == 0 {
514                        format!("{minutes:02}:{seconds:02}")
515                    } else {
516                        format!("{hours}:{minutes:02}:{seconds:02}")
517                    };
518                    Text::styled(formatted, style)
519                }
520            }
521        };
522        self.cache.borrow_mut().insert(task.id, (now, text.clone()));
523        text
524    }
525}
526
527/// `TaskProgressColumn.render_speed`: iterations per second with a power-of-ten
528/// suffix, e.g. `2.5×10³ it/s`.
529fn render_speed(speed: Option<f64>) -> Text {
530    let Some(speed) = speed else {
531        return Text::styled("", "progress.percentage");
532    };
533    let (unit, suffix) = filesize::pick_unit_and_suffix_signed(
534        speed as i64,
535        &["", "×10³", "×10⁶", "×10⁹", "×10¹²"],
536        1000,
537    );
538    let data_speed = speed / unit as f64;
539    Text::styled(
540        format!("{data_speed:.1}{suffix} it/s"),
541        "progress.percentage",
542    )
543}
544
545/// Python's `str(timedelta(seconds=n))` for `n >= 0`: `H:MM:SS`, prefixed by
546/// `N day(s), ` from a day upward.
547fn timedelta(total_seconds: i64) -> String {
548    let days = total_seconds / 86_400;
549    let rest = total_seconds % 86_400;
550    let clock = format!("{}:{:02}:{:02}", rest / 3600, rest % 3600 / 60, rest % 60);
551    match days {
552        0 => clock,
553        1 => format!("1 day, {clock}"),
554        days => format!("{days} days, {clock}"),
555    }
556}
557
558/// Python's `f"{value:,.{precision}f}"`: fixed precision with `,` grouping.
559fn grouped(value: f64, precision: usize) -> String {
560    let formatted = format!("{value:.precision$}");
561    let (sign, digits) = match formatted.strip_prefix('-') {
562        Some(rest) => ("-", rest),
563        None => ("", formatted.as_str()),
564    };
565    let (integer, fraction) = match digits.split_once('.') {
566        Some((integer, fraction)) => (integer, Some(fraction)),
567        None => (digits, None),
568    };
569    let mut grouped = String::new();
570    for (index, digit) in integer.chars().enumerate() {
571        if index > 0 && (integer.len() - index) % 3 == 0 {
572            grouped.push(',');
573        }
574        grouped.push(digit);
575    }
576    match fraction {
577        Some(fraction) => format!("{sign}{grouped}.{fraction}"),
578        None => format!("{sign}{grouped}"),
579    }
580}
581
582/// A single tracked task. Mirrors `rich.progress.Task`; read-only outside
583/// [`Progress`].
584pub struct Task {
585    id: TaskId,
586    description: String,
587    total: Option<f64>,
588    completed: f64,
589    visible: bool,
590    start_time: Option<f64>,
591    stop_time: Option<f64>,
592    finished_time: Option<f64>,
593    finished_speed: Option<f64>,
594    /// `(timestamp, completed)` speed samples (upstream `ProgressSample`).
595    samples: VecDeque<(f64, f64)>,
596    /// Arbitrary per-task values for format strings (upstream `fields`).
597    fields: BTreeMap<String, FormatValue>,
598    get_time: GetTime,
599}
600
601impl Task {
602    /// This task's custom fields (upstream `Task.fields`).
603    pub fn fields(&self) -> &BTreeMap<String, FormatValue> {
604        &self.fields
605    }
606
607    /// Resolve a `str.format` field name against this task, as
608    /// `text_format.format(task=task)` does: `task.<attribute>` or
609    /// `task.fields[<name>]`.
610    fn format_field(&self, name: &str) -> Option<FormatValue> {
611        let attribute = name.strip_prefix("task.")?;
612        if let Some(key) = attribute
613            .strip_prefix("fields[")
614            .and_then(|rest| rest.strip_suffix(']'))
615        {
616            return self.fields.get(key).cloned();
617        }
618        Some(match attribute {
619            "id" => FormatValue::Int(self.id.0 as i64),
620            "description" => FormatValue::Str(self.description.clone()),
621            // Python keeps the ints a caller passes (`total=200` prints `200`);
622            // this API takes floats, so a whole number formats as an int.
623            "total" => self.total.map_or(FormatValue::None, whole_number),
624            "completed" => whole_number(self.completed),
625            "visible" => FormatValue::Bool(self.visible),
626            "started" => FormatValue::Bool(self.started()),
627            "finished" => FormatValue::Bool(self.finished()),
628            "percentage" => FormatValue::Float(self.percentage()),
629            "remaining" => self.remaining().into(),
630            "elapsed" => self.elapsed().into(),
631            "speed" => self.speed().into(),
632            "time_remaining" => self.time_remaining().into(),
633            "start_time" => self.start_time.into(),
634            "stop_time" => self.stop_time.into(),
635            "finished_time" => self.finished_time.into(),
636            "finished_speed" => self.finished_speed.into(),
637            _ => return None,
638        })
639    }
640    fn now(&self) -> f64 {
641        (self.get_time)()
642    }
643
644    /// The current time by the progress's clock. Port of `Task.get_time`.
645    pub fn get_time(&self) -> f64 {
646        self.now()
647    }
648
649    /// This task's id.
650    pub fn id(&self) -> TaskId {
651        self.id
652    }
653
654    /// The description (console markup).
655    pub fn description(&self) -> &str {
656        &self.description
657    }
658
659    /// The total number of steps, or `None` when indeterminate.
660    pub fn total(&self) -> Option<f64> {
661        self.total
662    }
663
664    /// The number of steps completed.
665    pub fn completed(&self) -> f64 {
666        self.completed
667    }
668
669    /// Whether the task is shown.
670    pub fn visible(&self) -> bool {
671        self.visible
672    }
673
674    /// Whether the task has been started.
675    pub fn started(&self) -> bool {
676        self.start_time.is_some()
677    }
678
679    /// Steps left, or `None` when indeterminate.
680    pub fn remaining(&self) -> Option<f64> {
681        self.total.map(|total| total - self.completed)
682    }
683
684    /// Seconds since the task started (to its stop time, if stopped).
685    pub fn elapsed(&self) -> Option<f64> {
686        let start = self.start_time?;
687        Some(self.stop_time.unwrap_or_else(|| self.now()) - start)
688    }
689
690    /// Whether the task has reached its total.
691    pub fn finished(&self) -> bool {
692        self.finished_time.is_some()
693    }
694
695    /// The elapsed time recorded when the task finished.
696    pub fn finished_time(&self) -> Option<f64> {
697        self.finished_time
698    }
699
700    /// The completion percentage, clamped to 0–100 (0 without a total).
701    pub fn percentage(&self) -> f64 {
702        match self.total {
703            Some(total) if total != 0.0 => (self.completed / total * 100.0).clamp(0.0, 100.0),
704            _ => 0.0,
705        }
706    }
707
708    /// Steps per second over the sample window, or `None` without enough samples.
709    pub fn speed(&self) -> Option<f64> {
710        self.start_time?;
711        let (first, _) = *self.samples.front()?;
712        let (last, _) = *self.samples.back()?;
713        let total_time = last - first;
714        if total_time == 0.0 {
715            return None;
716        }
717        let total_completed: f64 = self.samples.iter().skip(1).map(|(_, done)| done).sum();
718        Some(total_completed / total_time)
719    }
720
721    /// Estimated seconds remaining (rounded up), 0 once finished.
722    pub fn time_remaining(&self) -> Option<f64> {
723        if self.finished() {
724            return Some(0.0);
725        }
726        let speed = self.speed().filter(|speed| *speed != 0.0)?;
727        let remaining = self.remaining()?;
728        Some((remaining / speed).ceil())
729    }
730
731    /// `Task._reset`.
732    fn clear_progress(&mut self) {
733        self.samples.clear();
734        self.finished_time = None;
735        self.finished_speed = None;
736    }
737
738    /// The percentage cell: `[progress.percentage]{percentage:>3.0f}%`, empty
739    /// without a total (`text_format_no_percentage`).
740    fn percentage_cell(&self) -> Text {
741        if self.total.is_none() {
742            return Text::new("");
743        }
744        let mut text = Text::new(format!("{:>3.0}%", self.percentage()));
745        let len = text.plain().len();
746        text.stylize("progress.percentage", 0, len);
747        text
748    }
749
750    /// The M-of-N cell text: `completed` right-justified to the width of `total`
751    /// (`?` when indeterminate), then `/total`. Port of `MofNCompleteColumn.render`.
752    fn mofn_text(&self) -> String {
753        let completed = self.completed as i64;
754        let total = self
755            .total
756            .map_or_else(|| "?".to_string(), |total| (total as i64).to_string());
757        let total_width = total.chars().count();
758        format!("{completed:>total_width$}/{total}")
759    }
760
761    /// The download cell text: `completed`/`total` in a shared byte unit, e.g.
762    /// `0.5/1.0 kB`. Port of `DownloadColumn.render`.
763    fn download_text(&self, binary: bool) -> String {
764        const DECIMAL: &[&str] = &["bytes", "kB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"];
765        const BINARY: &[&str] = &[
766            "bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB",
767        ];
768        // `int(task.completed)` / `int(task.total)`: truncated, sign kept.
769        let completed = self.completed as i64;
770        let base_size = self.total.map_or(completed, |total| total as i64);
771        let (unit, suffix) = if binary {
772            filesize::pick_unit_and_suffix_signed(base_size, BINARY, 1024)
773        } else {
774            filesize::pick_unit_and_suffix_signed(base_size, DECIMAL, 1000)
775        };
776        let precision = if unit == 1 { 0 } else { 1 };
777        let completed_str = grouped(completed as f64 / unit as f64, precision);
778        let total_str = self.total.map_or_else(
779            || "?".to_string(),
780            |total| grouped((total as i64) as f64 / unit as f64, precision),
781        );
782        format!("{completed_str}/{total_str} {suffix}")
783    }
784}
785
786/// Changes for [`Progress::update`]; unset fields are left alone. Upstream's
787/// keyword arguments to `Progress.update`.
788#[derive(Debug, Clone, Default)]
789pub struct TaskUpdate {
790    pub total: Option<f64>,
791    pub completed: Option<f64>,
792    pub advance: Option<f64>,
793    pub description: Option<String>,
794    pub visible: Option<bool>,
795    /// Custom fields to set (upstream `**fields`).
796    pub fields: Vec<(String, FormatValue)>,
797    /// Redraw a live display right after the update (upstream `refresh=True`).
798    pub refresh: bool,
799}
800
801impl TaskUpdate {
802    pub fn total(mut self, total: f64) -> Self {
803        self.total = Some(total);
804        self
805    }
806
807    pub fn completed(mut self, completed: f64) -> Self {
808        self.completed = Some(completed);
809        self
810    }
811
812    pub fn advance(mut self, advance: f64) -> Self {
813        self.advance = Some(advance);
814        self
815    }
816
817    pub fn description(mut self, description: impl Into<String>) -> Self {
818        self.description = Some(description.into());
819        self
820    }
821
822    pub fn visible(mut self, visible: bool) -> Self {
823        self.visible = Some(visible);
824        self
825    }
826
827    /// Redraw a [`LiveProgress`] right after this update (upstream
828    /// `refresh=True`).
829    pub fn refresh(mut self, refresh: bool) -> Self {
830        self.refresh = refresh;
831        self
832    }
833
834    /// Set a custom field (upstream `update(task_id, **fields)`).
835    pub fn field(mut self, name: impl Into<String>, value: impl Into<FormatValue>) -> Self {
836        self.fields.push((name.into(), value.into()));
837        self
838    }
839}
840
841/// A progress display over one or more [`Task`]s. Mirrors `rich.progress.Progress`.
842pub struct Progress {
843    tasks: Vec<Task>,
844    next_id: usize,
845    columns: Vec<ProgressColumn>,
846    get_time: GetTime,
847    speed_estimate_period: f64,
848    expand: bool,
849    transient: bool,
850    disable: bool,
851}
852
853impl Default for Progress {
854    fn default() -> Self {
855        Progress {
856            tasks: Vec::new(),
857            next_id: 0,
858            columns: Progress::default_columns(),
859            get_time: Arc::new(monotonic),
860            speed_estimate_period: 30.0,
861            expand: false,
862            transient: false,
863            disable: false,
864        }
865    }
866}
867
868impl Progress {
869    pub fn new() -> Self {
870        Progress::default()
871    }
872
873    /// Upstream's `Progress.get_default_columns()`: description, bar,
874    /// percentage and time remaining.
875    pub fn default_columns() -> Vec<ProgressColumn> {
876        vec![
877            ProgressColumn::Description,
878            ProgressColumn::Bar,
879            ProgressColumn::Percentage,
880            ProgressColumn::time_remaining(),
881        ]
882    }
883
884    /// Replace the column list (default: [`default_columns`](Self::default_columns)).
885    pub fn columns(mut self, columns: Vec<ProgressColumn>) -> Self {
886        self.columns = columns;
887        self
888    }
889
890    /// Read time from `clock` (seconds) instead of the monotonic clock.
891    /// Upstream's `get_time`; makes time-based columns deterministic.
892    pub fn clock(mut self, clock: impl Fn() -> f64 + Send + Sync + 'static) -> Self {
893        self.get_time = Arc::new(clock);
894        for task in &mut self.tasks {
895            task.get_time = self.get_time.clone();
896        }
897        self
898    }
899
900    /// Stretch the task grid to the full width (upstream `expand`).
901    pub fn expand(mut self, expand: bool) -> Self {
902        self.expand = expand;
903        self
904    }
905
906    /// Erase the display when it stops (upstream `transient`).
907    pub fn transient(mut self, transient: bool) -> Self {
908        self.transient = transient;
909        self
910    }
911
912    /// Show nothing: [`start`](Self::start) draws no display, while tasks
913    /// still update (upstream `disable`).
914    pub fn disable(mut self, disable: bool) -> Self {
915        self.disable = disable;
916        self
917    }
918
919    /// Seconds of history used for speed estimates (default 30).
920    pub fn speed_estimate_period(mut self, seconds: f64) -> Self {
921        self.speed_estimate_period = seconds;
922        self
923    }
924
925    fn now(&self) -> f64 {
926        (self.get_time)()
927    }
928
929    fn task_mut(&mut self, id: TaskId) -> Option<&mut Task> {
930        self.tasks.iter_mut().find(|task| task.id == id)
931    }
932
933    /// Add a started task and return its id. Port of `Progress.add_task`
934    /// (`start=True`); `total` of `None` is an indeterminate task.
935    pub fn add_task(
936        &mut self,
937        description: impl Into<String>,
938        total: impl Into<Option<f64>>,
939        completed: f64,
940    ) -> TaskId {
941        let id = self.push_task(description.into(), total.into(), completed);
942        self.start_task(id);
943        id
944    }
945
946    /// Add a task that has not started (`add_task(start=False)`): it shows no
947    /// elapsed time until [`start_task`](Self::start_task).
948    pub fn add_unstarted_task(
949        &mut self,
950        description: impl Into<String>,
951        total: impl Into<Option<f64>>,
952        completed: f64,
953    ) -> TaskId {
954        self.push_task(description.into(), total.into(), completed)
955    }
956
957    /// Add a task with custom fields for format strings. Port of
958    /// `add_task(description, total=…, completed=…, start=…, **fields)`.
959    pub fn add_task_with<K: Into<String>, V: Into<FormatValue>>(
960        &mut self,
961        description: impl Into<String>,
962        total: impl Into<Option<f64>>,
963        completed: f64,
964        start: bool,
965        fields: impl IntoIterator<Item = (K, V)>,
966    ) -> TaskId {
967        let id = self.push_task(description.into(), total.into(), completed);
968        if let Some(task) = self.task_mut(id) {
969            task.fields = fields
970                .into_iter()
971                .map(|(name, value)| (name.into(), value.into()))
972                .collect();
973        }
974        if start {
975            self.start_task(id);
976        }
977        id
978    }
979
980    fn push_task(&mut self, description: String, total: Option<f64>, completed: f64) -> TaskId {
981        let id = TaskId(self.next_id);
982        self.next_id += 1;
983        self.tasks.push(Task {
984            id,
985            description,
986            total,
987            completed,
988            visible: true,
989            start_time: None,
990            stop_time: None,
991            finished_time: None,
992            finished_speed: None,
993            samples: VecDeque::new(),
994            fields: BTreeMap::new(),
995            get_time: self.get_time.clone(),
996        });
997        id
998    }
999
1000    /// The task with this id, if it has not been removed.
1001    pub fn task(&self, id: TaskId) -> Option<&Task> {
1002        self.tasks.iter().find(|task| task.id == id)
1003    }
1004
1005    /// Every task, in the order added.
1006    pub fn tasks(&self) -> &[Task] {
1007        &self.tasks
1008    }
1009
1010    /// Whether every task has finished. Port of `Progress.finished`.
1011    pub fn finished(&self) -> bool {
1012        self.tasks.iter().all(Task::finished)
1013    }
1014
1015    /// Start a task's clock if it has not started. Port of `start_task`.
1016    pub fn start_task(&mut self, id: TaskId) {
1017        let now = self.now();
1018        if let Some(task) = self.task_mut(id) {
1019            task.start_time.get_or_insert(now);
1020        }
1021    }
1022
1023    /// Stop a task's clock; its elapsed time freezes. Port of `stop_task`.
1024    pub fn stop_task(&mut self, id: TaskId) {
1025        let now = self.now();
1026        if let Some(task) = self.task_mut(id) {
1027            task.start_time.get_or_insert(now);
1028            task.stop_time = Some(now);
1029        }
1030    }
1031
1032    /// Update a task. Port of `Progress.update`: a new total clears the speed
1033    /// samples; positive progress adds a sample; reaching the total records the
1034    /// finish time.
1035    pub fn update(&mut self, id: TaskId, update: TaskUpdate) {
1036        let now = self.now();
1037        let period = self.speed_estimate_period;
1038        let Some(task) = self.task_mut(id) else {
1039            return;
1040        };
1041        let completed_start = task.completed;
1042        if let Some(total) = update.total {
1043            if Some(total) != task.total {
1044                task.total = Some(total);
1045                task.clear_progress();
1046            }
1047        }
1048        if let Some(advance) = update.advance {
1049            task.completed += advance;
1050        }
1051        if let Some(completed) = update.completed {
1052            task.completed = completed;
1053        }
1054        if let Some(description) = update.description {
1055            task.description = description;
1056        }
1057        if let Some(visible) = update.visible {
1058            task.visible = visible;
1059        }
1060        task.fields.extend(update.fields);
1061        let update_completed = task.completed - completed_start;
1062        let old_sample_time = now - period;
1063        while task
1064            .samples
1065            .front()
1066            .is_some_and(|(time, _)| *time < old_sample_time)
1067        {
1068            task.samples.pop_front();
1069        }
1070        if update_completed > 0.0 {
1071            task.samples.push_back((now, update_completed));
1072            if task.samples.len() > MAX_SAMPLES {
1073                task.samples.pop_front();
1074            }
1075        }
1076        if task.total.is_some_and(|total| task.completed >= total) && task.finished_time.is_none() {
1077            task.finished_time = task.elapsed();
1078        }
1079    }
1080
1081    /// Advance a task by `amount` steps. Port of `Progress.advance`, which
1082    /// (unlike `update`) always records a sample and the finish speed.
1083    pub fn advance(&mut self, id: TaskId, amount: f64) {
1084        let now = self.now();
1085        let period = self.speed_estimate_period;
1086        let Some(task) = self.task_mut(id) else {
1087            return;
1088        };
1089        let completed_start = task.completed;
1090        task.completed += amount;
1091        let update_completed = task.completed - completed_start;
1092        let old_sample_time = now - period;
1093        while task
1094            .samples
1095            .front()
1096            .is_some_and(|(time, _)| *time < old_sample_time)
1097        {
1098            task.samples.pop_front();
1099        }
1100        while task.samples.len() > MAX_SAMPLES {
1101            task.samples.pop_front();
1102        }
1103        task.samples.push_back((now, update_completed));
1104        if task.samples.len() > MAX_SAMPLES {
1105            task.samples.pop_front();
1106        }
1107        if task.total.is_some_and(|total| task.completed >= total) && task.finished_time.is_none() {
1108            task.finished_time = task.elapsed();
1109            task.finished_speed = task.speed();
1110        }
1111    }
1112
1113    /// Reset a task to `completed`, optionally restarting its clock and
1114    /// changing its total. Port of `Progress.reset`. Like upstream, a stop
1115    /// time set earlier is kept.
1116    pub fn reset(&mut self, id: TaskId, start: bool, total: Option<f64>, completed: f64) {
1117        let now = self.now();
1118        let Some(task) = self.task_mut(id) else {
1119            return;
1120        };
1121        task.clear_progress();
1122        task.start_time = start.then_some(now);
1123        if let Some(total) = total {
1124            task.total = Some(total);
1125        }
1126        task.completed = completed;
1127        task.finished_time = None;
1128    }
1129
1130    /// Remove a task. Port of `Progress.remove_task`.
1131    pub fn remove_task(&mut self, id: TaskId) {
1132        self.tasks.retain(|task| task.id != id);
1133    }
1134}
1135
1136impl Progress {
1137    /// The grid the display renders. Port of `Progress.make_tasks_table`: a
1138    /// `Table.grid` with one column per [`ProgressColumn`] (its table column
1139    /// options), `padding=(0, 1)` and the progress's `expand`, and one row per
1140    /// visible task.
1141    pub fn make_tasks_table(&self) -> Table {
1142        let mut table = Table::grid().padding(0, 1, 0, 1).expand(self.expand);
1143        for column in &self.columns {
1144            table.add_column_with(Text::new(""), column.table_column());
1145        }
1146        for task in self.tasks.iter().filter(|task| task.visible) {
1147            // Each column is called once per row, in order: spinners and the
1148            // remaining-time cache are stateful, as upstream's columns are.
1149            let cells = self
1150                .columns
1151                .iter()
1152                .map(|column| column.table_cell(task))
1153                .collect();
1154            table.add_row_cells(cells);
1155        }
1156        table
1157    }
1158}
1159
1160impl Renderable for Progress {
1161    fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
1162        self.make_tasks_table().rich_render(console, options)
1163    }
1164}
1165
1166/// A [`Progress`] shared with the auto-refresh thread of a live display.
1167struct ProgressView(Arc<std::sync::Mutex<Progress>>);
1168
1169impl Renderable for ProgressView {
1170    fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
1171        match self.0.lock() {
1172            Ok(progress) => progress.rich_render(console, options),
1173            Err(poisoned) => poisoned.into_inner().rich_render(console, options),
1174        }
1175    }
1176}
1177
1178impl Progress {
1179    /// Start an auto-refreshing live display of this progress, redrawn
1180    /// `refresh_per_second` times a second on a background thread. Port of
1181    /// `Progress.start` with `auto_refresh=True`; stop it with
1182    /// [`LiveProgress::stop`] (upstream's `with progress:` block).
1183    pub fn start<W: std::io::Write + Send + 'static>(
1184        self,
1185        console: Console,
1186        writer: W,
1187        refresh_per_second: f64,
1188    ) -> LiveProgress<W> {
1189        // `disable` draws nothing: upstream skips `live.start()` and `stop()`.
1190        if self.disable {
1191            return LiveProgress {
1192                progress: Arc::new(std::sync::Mutex::new(self)),
1193                live: None,
1194                writer: Some(writer),
1195                interactive: true,
1196                holder: std::sync::Mutex::new(None),
1197            };
1198        }
1199        let transient = self.transient;
1200        let interactive = console.is_terminal();
1201        let shared = Arc::new(std::sync::Mutex::new(self));
1202        let live = crate::live::Live::spawn_with(
1203            Box::new(ProgressView(shared.clone())),
1204            console,
1205            writer,
1206            refresh_per_second,
1207            transient,
1208        );
1209        LiveProgress {
1210            progress: shared,
1211            live: Some(live),
1212            writer: None,
1213            interactive,
1214            holder: std::sync::Mutex::new(None),
1215        }
1216    }
1217}
1218
1219/// A running [`Progress`] display. Task changes made through it are picked up
1220/// by the next refresh, as with upstream's `refresh=False` updates.
1221pub struct LiveProgress<W: std::io::Write + Send + 'static> {
1222    progress: Arc<std::sync::Mutex<Progress>>,
1223    live: Option<crate::live::AutoLive<W>>,
1224    /// The sink of a disabled display, which never reaches a live thread.
1225    writer: Option<W>,
1226    /// Whether the console is a terminal; `stop` ends a file with a newline.
1227    interactive: bool,
1228    /// The thread inside [`with`](LiveProgress::with), if any. Upstream's lock
1229    /// is an `RLock`; ours is not, so re-entry is detected rather than left to
1230    /// deadlock on the mutex or on the refresh thread.
1231    holder: std::sync::Mutex<Option<std::thread::ThreadId>>,
1232}
1233
1234/// Clears [`LiveProgress`]'s `holder` when a `with` block ends, unwinding
1235/// included.
1236struct HolderGuard<'a>(&'a std::sync::Mutex<Option<std::thread::ThreadId>>);
1237
1238impl Drop for HolderGuard<'_> {
1239    fn drop(&mut self) {
1240        *self
1241            .0
1242            .lock()
1243            .unwrap_or_else(|poisoned| poisoned.into_inner()) = None;
1244    }
1245}
1246
1247impl<W: std::io::Write + Send + 'static> LiveProgress<W> {
1248    /// Run `f` with the progress locked, for any change not wrapped below.
1249    ///
1250    /// `f` may call [`refresh`](Self::refresh) (the frame is drawn as soon as
1251    /// the lock is released), but not `with` or a method built on it: upstream
1252    /// re-enters its `RLock`, which a `&mut Progress` cannot express, so a
1253    /// nested call panics instead of deadlocking.
1254    pub fn with<R>(&self, f: impl FnOnce(&mut Progress) -> R) -> R {
1255        let current = std::thread::current().id();
1256        assert!(
1257            !self.held_by(current),
1258            "LiveProgress::with re-entered from inside a `with` closure; \
1259             use the `&mut Progress` it was given instead"
1260        );
1261        let mut progress = match self.progress.lock() {
1262            Ok(progress) => progress,
1263            Err(poisoned) => poisoned.into_inner(),
1264        };
1265        *self
1266            .holder
1267            .lock()
1268            .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(current);
1269        let _holder = HolderGuard(&self.holder);
1270        f(&mut progress)
1271    }
1272
1273    /// Whether `thread` is inside [`with`](Self::with) right now.
1274    fn held_by(&self, thread: std::thread::ThreadId) -> bool {
1275        *self
1276            .holder
1277            .lock()
1278            .unwrap_or_else(|poisoned| poisoned.into_inner())
1279            == Some(thread)
1280    }
1281
1282    /// [`Progress::add_task`], then redraw (upstream's `add_task` refreshes).
1283    pub fn add_task(
1284        &self,
1285        description: impl Into<String>,
1286        total: impl Into<Option<f64>>,
1287        completed: f64,
1288    ) -> TaskId {
1289        let id = self.with(|progress| progress.add_task(description, total, completed));
1290        self.refresh();
1291        id
1292    }
1293
1294    /// [`Progress::reset`], then redraw (upstream's `reset` refreshes).
1295    pub fn reset(&self, id: TaskId, start: bool, total: Option<f64>, completed: f64) {
1296        self.with(|progress| progress.reset(id, start, total, completed));
1297        self.refresh();
1298    }
1299
1300    /// [`Progress::advance`].
1301    pub fn advance(&self, id: TaskId, amount: f64) {
1302        self.with(|progress| progress.advance(id, amount));
1303    }
1304
1305    /// [`Progress::update`]; redraws when the update asks to
1306    /// ([`TaskUpdate::refresh`]).
1307    pub fn update(&self, id: TaskId, update: TaskUpdate) {
1308        let refresh = update.refresh;
1309        self.with(|progress| progress.update(id, update));
1310        if refresh {
1311            self.refresh();
1312        }
1313    }
1314
1315    /// Redraw now rather than at the next tick, returning once the frame is
1316    /// written. Port of `Progress.refresh`.
1317    ///
1318    /// Called from inside [`with`](Self::with), the redraw is queued instead:
1319    /// the refresh thread needs the lock this thread holds to render, so
1320    /// waiting for it would deadlock. The frame is drawn when `with` returns.
1321    pub fn refresh(&self) {
1322        if let Some(live) = &self.live {
1323            if self.held_by(std::thread::current().id()) {
1324                live.refresh();
1325            } else {
1326                live.refresh_wait();
1327            }
1328        }
1329    }
1330
1331    /// Iterate `iter`, advancing a new task by one after each item is
1332    /// processed. Port of `Progress.track`: the total defaults to the
1333    /// iterator's exact length, else the task is indeterminate.
1334    pub fn track<I: IntoIterator>(
1335        &self,
1336        iter: I,
1337        total: Option<f64>,
1338        description: impl Into<String>,
1339    ) -> Track<'_, I::IntoIter, W> {
1340        let iter = iter.into_iter();
1341        let total = total.or_else(|| match iter.size_hint() {
1342            (lower, Some(upper)) if lower == upper && lower > 0 => Some(lower as f64),
1343            _ => None,
1344        });
1345        let task = self.add_task(description, total, 0.0);
1346        Track {
1347            iter,
1348            progress: self,
1349            task,
1350            pending: false,
1351        }
1352    }
1353
1354    /// Track reading from `reader`: each read advances the task by the bytes
1355    /// read. Port of `Progress.wrap_file`: `total` is the byte count, or else
1356    /// the total of `task`; a new task named `description` is added when
1357    /// `task` is `None`, otherwise `task`'s total is set.
1358    pub fn wrap_read<R: std::io::Read>(
1359        &self,
1360        reader: R,
1361        total: Option<u64>,
1362        task: Option<TaskId>,
1363        description: impl Into<String>,
1364    ) -> std::io::Result<ProgressReader<'_, R, W>> {
1365        let total = total.map(|total| total as f64).or_else(|| {
1366            task.and_then(|task| self.with(|progress| progress.task(task).and_then(Task::total)))
1367        });
1368        let Some(total) = total else {
1369            return Err(std::io::Error::new(
1370                std::io::ErrorKind::InvalidInput,
1371                "unable to get the total number of bytes, please specify 'total'",
1372            ));
1373        };
1374        let task = self.task_for(task, total, description);
1375        Ok(ProgressReader {
1376            reader,
1377            progress: self,
1378            task,
1379        })
1380    }
1381
1382    /// Open `path` for reading and track it. Port of `Progress.open` in
1383    /// binary mode: `total` defaults to the file's size.
1384    pub fn open(
1385        &self,
1386        path: impl AsRef<std::path::Path>,
1387        total: Option<u64>,
1388        task: Option<TaskId>,
1389        description: impl Into<String>,
1390    ) -> std::io::Result<ProgressReader<'_, std::fs::File, W>> {
1391        let file = std::fs::File::open(path)?;
1392        let total = match total {
1393            Some(total) => total,
1394            None => file.metadata()?.len(),
1395        };
1396        let task = self.task_for(task, total as f64, description);
1397        Ok(ProgressReader {
1398            reader: file,
1399            progress: self,
1400            task,
1401        })
1402    }
1403
1404    /// A new task with `total`, or `task` with its total set to it.
1405    fn task_for(&self, task: Option<TaskId>, total: f64, description: impl Into<String>) -> TaskId {
1406        match task {
1407            Some(task) => {
1408                self.update(task, TaskUpdate::default().total(total));
1409                task
1410            }
1411            None => self.add_task(description, total, 0.0),
1412        }
1413    }
1414
1415    /// Commit the final frame, stop the refresh thread, and return the
1416    /// progress and the output sink. Port of `Progress.stop`.
1417    pub fn stop(mut self) -> (Progress, W) {
1418        let writer = match self.live.take() {
1419            Some(live) => {
1420                let mut writer = live.stop();
1421                // `Progress.stop`: `console.print()` when not interactive.
1422                if !self.interactive {
1423                    let _ = writer.write_all(b"\n");
1424                }
1425                writer
1426            }
1427            None => self
1428                .writer
1429                .take()
1430                .expect("a disabled display keeps its writer"),
1431        };
1432        let progress = match Arc::try_unwrap(std::mem::replace(
1433            &mut self.progress,
1434            Arc::new(std::sync::Mutex::new(Progress::new())),
1435        )) {
1436            Ok(mutex) => mutex
1437                .into_inner()
1438                .unwrap_or_else(|poisoned| poisoned.into_inner()),
1439            // The refresh thread has exited, so this is the only owner left.
1440            Err(_) => unreachable!("progress still shared after the live display stopped"),
1441        };
1442        (progress, writer)
1443    }
1444}
1445
1446/// A reader that advances a task by the bytes read through it. Returned by
1447/// [`LiveProgress::wrap_read`] and [`LiveProgress::open`] (upstream `_Reader`).
1448pub struct ProgressReader<'a, R, W: std::io::Write + Send + 'static> {
1449    reader: R,
1450    progress: &'a LiveProgress<W>,
1451    task: TaskId,
1452}
1453
1454impl<R, W: std::io::Write + Send + 'static> ProgressReader<'_, R, W> {
1455    /// The task this reader advances.
1456    pub fn task(&self) -> TaskId {
1457        self.task
1458    }
1459
1460    /// The wrapped reader.
1461    pub fn into_inner(self) -> R {
1462        self.reader
1463    }
1464}
1465
1466impl<R: std::io::Read, W: std::io::Write + Send + 'static> std::io::Read
1467    for ProgressReader<'_, R, W>
1468{
1469    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
1470        let count = self.reader.read(buf)?;
1471        self.progress.advance(self.task, count as f64);
1472        Ok(count)
1473    }
1474}
1475
1476impl<R: std::io::BufRead, W: std::io::Write + Send + 'static> std::io::BufRead
1477    for ProgressReader<'_, R, W>
1478{
1479    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
1480        self.reader.fill_buf()
1481    }
1482
1483    fn consume(&mut self, amount: usize) {
1484        self.reader.consume(amount);
1485        self.progress.advance(self.task, amount as f64);
1486    }
1487}
1488
1489/// The iterator [`LiveProgress::track`] returns.
1490pub struct Track<'a, I: Iterator, W: std::io::Write + Send + 'static> {
1491    iter: I,
1492    progress: &'a LiveProgress<W>,
1493    task: TaskId,
1494    /// Whether an item has been handed out but not yet counted: upstream
1495    /// advances after the loop body, when the next item is requested.
1496    pending: bool,
1497}
1498
1499impl<I: Iterator, W: std::io::Write + Send + 'static> Track<'_, I, W> {
1500    /// The task this iterator advances.
1501    pub fn task(&self) -> TaskId {
1502        self.task
1503    }
1504}
1505
1506impl<I: Iterator, W: std::io::Write + Send + 'static> Iterator for Track<'_, I, W> {
1507    type Item = I::Item;
1508
1509    fn next(&mut self) -> Option<I::Item> {
1510        if std::mem::take(&mut self.pending) {
1511            self.progress.advance(self.task, 1.0);
1512        }
1513        let item = self.iter.next();
1514        if item.is_some() {
1515            self.pending = true;
1516        } else {
1517            self.progress.refresh();
1518        }
1519        item
1520    }
1521
1522    fn size_hint(&self) -> (usize, Option<usize>) {
1523        self.iter.size_hint()
1524    }
1525}
1526
1527/// Track progress over `iter` with a live display on stdout. Port of the
1528/// module-level `rich.progress.track`: the description, bar, progress and
1529/// time-remaining columns, refreshed ten times a second, stopped when the
1530/// iterator is exhausted or dropped.
1531pub fn track<I: IntoIterator>(iter: I, description: &str) -> TrackStdout<I::IntoIter> {
1532    let mut columns = Vec::new();
1533    if !description.is_empty() {
1534        columns.push(ProgressColumn::Description);
1535    }
1536    columns.extend([
1537        ProgressColumn::Bar,
1538        ProgressColumn::TaskProgress { show_speed: true },
1539        ProgressColumn::TimeRemaining(TimeRemainingColumn::new(false, true)),
1540    ]);
1541    let iter = iter.into_iter();
1542    let total = match iter.size_hint() {
1543        (lower, Some(upper)) if lower == upper && lower > 0 => Some(lower as f64),
1544        _ => None,
1545    };
1546    let live = Progress::new()
1547        .columns(columns)
1548        .start(Console::new(), std::io::stdout(), 10.0);
1549    let task = live.add_task(description, total, 0.0);
1550    TrackStdout {
1551        iter,
1552        live: Some(live),
1553        task,
1554        pending: false,
1555    }
1556}
1557
1558/// The iterator [`track`] returns; it owns its live display.
1559pub struct TrackStdout<I: Iterator> {
1560    iter: I,
1561    live: Option<LiveProgress<std::io::Stdout>>,
1562    task: TaskId,
1563    pending: bool,
1564}
1565
1566impl<I: Iterator> Iterator for TrackStdout<I> {
1567    type Item = I::Item;
1568
1569    fn next(&mut self) -> Option<I::Item> {
1570        let live = self.live.as_ref()?;
1571        if std::mem::take(&mut self.pending) {
1572            live.advance(self.task, 1.0);
1573        }
1574        match self.iter.next() {
1575            Some(item) => {
1576                self.pending = true;
1577                Some(item)
1578            }
1579            None => {
1580                if let Some(live) = self.live.take() {
1581                    live.stop();
1582                }
1583                None
1584            }
1585        }
1586    }
1587}
1588
1589impl<I: Iterator> Drop for TrackStdout<I> {
1590    fn drop(&mut self) {
1591        if let Some(live) = self.live.take() {
1592            live.stop();
1593        }
1594    }
1595}
1596
1597/// A whole, exactly representable count as an int, else the float.
1598fn whole_number(value: f64) -> FormatValue {
1599    if value.fract() == 0.0 && value.abs() < 9_007_199_254_740_992.0 {
1600        FormatValue::Int(value as i64)
1601    } else {
1602        FormatValue::Float(value)
1603    }
1604}
1605
1606#[cfg(test)]
1607mod tests {
1608    use super::*;
1609    use crate::color::ColorSystem;
1610
1611    fn render(progress: &Progress) -> String {
1612        Console::builder()
1613            .force_terminal(true)
1614            .color_system(Some(ColorSystem::Truecolor))
1615            .width(50)
1616            .no_color(false)
1617            .build()
1618            .render_to_string(progress)
1619    }
1620
1621    #[test]
1622    fn three_tasks_match_upstream() {
1623        // Captured from real rich 15.0.0 (default columns, width 50).
1624        let mut progress = Progress::new().columns(vec![
1625            ProgressColumn::Description,
1626            ProgressColumn::Bar,
1627            ProgressColumn::Percentage,
1628        ]);
1629        progress.add_task("Downloading", 100.0, 50.0);
1630        progress.add_task("Processing", 100.0, 100.0);
1631        progress.add_task("Waiting", 100.0, 0.0);
1632        let expected = concat!(
1633            "Downloading \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━\x1b[0m",
1634            "\x1b[38;2;249;38;114m╸\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 50%\x1b[0m\n",
1635            "Processing  \x1b[38;2;114;156;31m",
1636            "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m\n",
1637            "Waiting     \x1b[38;5;237m",
1638            "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m  0%\x1b[0m",
1639        );
1640        assert_eq!(render(&progress), expected);
1641    }
1642
1643    #[test]
1644    fn download_text_matches_upstream() {
1645        // Captured from real rich 15.0.0 DownloadColumn.render (decimal units).
1646        let dl = |completed: f64, total: f64| {
1647            let mut progress = Progress::new();
1648            let id = progress.add_task("", total, completed);
1649            progress.task(id).unwrap().download_text(false)
1650        };
1651        assert_eq!(dl(500.0, 1000.0), "0.5/1.0 kB");
1652        assert_eq!(dl(500.0, 999.0), "500/999 bytes");
1653        assert_eq!(dl(1_500_000.0, 3_000_000.0), "1.5/3.0 MB");
1654        assert_eq!(dl(0.0, 1024.0), "0.0/1.0 kB");
1655        assert_eq!(dl(2_500_000_000.0, 10_000_000_000.0), "2.5/10.0 GB");
1656        assert_eq!(dl(250.0, 250.0), "250/250 bytes");
1657    }
1658
1659    #[test]
1660    fn download_column_in_grid_matches_upstream() {
1661        // Captured from real rich 15.0.0: description + bar + download at width 50.
1662        let mut progress = Progress::new().columns(vec![
1663            ProgressColumn::Description,
1664            ProgressColumn::Bar,
1665            ProgressColumn::Download,
1666        ]);
1667        progress.add_task("File", 1000.0, 500.0);
1668        let expected = concat!(
1669            "File \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m",
1670            "\x1b[38;5;237m━━━━━━━━━━━━━━━━\x1b[0m \x1b[32m0.5/1.0 kB\x1b[0m",
1671        );
1672        assert_eq!(render(&progress), expected);
1673    }
1674
1675    #[test]
1676    fn custom_columns_with_mofn_match_upstream() {
1677        // Captured from real rich 15.0.0: description + bar + M-of-N (differing
1678        // M-of-N widths → the narrower cell left-justifies with green padding).
1679        let mut progress = Progress::new().columns(vec![
1680            ProgressColumn::Description,
1681            ProgressColumn::Bar,
1682            ProgressColumn::MofN,
1683        ]);
1684        progress.add_task("A", 5.0, 3.0);
1685        progress.add_task("B", 100.0, 50.0);
1686        let console = Console::builder()
1687            .force_terminal(true)
1688            .color_system(Some(ColorSystem::Truecolor))
1689            .width(40)
1690            .no_color(false)
1691            .build();
1692        let expected = concat!(
1693            "A \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m",
1694            "\x1b[38;5;237m━━━━━━━━━━━\x1b[0m \x1b[32m3/5    \x1b[0m\n",
1695            "B \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m",
1696            "\x1b[38;5;237m━━━━━━━━━━━━━━\x1b[0m \x1b[32m 50/100\x1b[0m",
1697        );
1698        assert_eq!(console.render_to_string(&progress), expected);
1699    }
1700
1701    fn live(columns: Vec<ProgressColumn>) -> LiveProgress<Vec<u8>> {
1702        let console = Console::builder()
1703            .force_terminal(true)
1704            .color_system(Some(ColorSystem::Truecolor))
1705            .width(40)
1706            .build();
1707        Progress::new()
1708            .columns(columns)
1709            .clock(|| 0.0)
1710            .start(console, Vec::new(), 1e-9)
1711    }
1712
1713    /// Run `f` on its own thread and fail (rather than hang the suite) if it
1714    /// has not finished within a few seconds.
1715    fn within_deadline<R: Send + 'static>(f: impl FnOnce() -> R + Send + 'static) -> R {
1716        let (done, wait) = std::sync::mpsc::channel();
1717        std::thread::spawn(move || {
1718            let _ = done.send(std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)));
1719        });
1720        match wait.recv_timeout(std::time::Duration::from_secs(10)) {
1721            Ok(Ok(value)) => value,
1722            Ok(Err(payload)) => std::panic::resume_unwind(payload),
1723            Err(_) => panic!("deadlocked: did not finish within 10s"),
1724        }
1725    }
1726
1727    #[test]
1728    fn refresh_inside_with_does_not_deadlock() {
1729        // Upstream's `Progress` lock is an `RLock` and `refresh()` renders in
1730        // the caller's thread, so refreshing while holding the lock is fine.
1731        let output = within_deadline(|| {
1732            let live = live(vec![ProgressColumn::Description, ProgressColumn::MofN]);
1733            live.with(|progress| {
1734                let task = progress.add_task("inside", Some(2.0), 1.0);
1735                live.refresh();
1736                task
1737            });
1738            live.refresh();
1739            String::from_utf8(live.stop().1).unwrap()
1740        });
1741        assert!(
1742            output.ends_with("inside \x1b[32m1/2\x1b[0m\n\x1b[?25h"),
1743            "{output:?}"
1744        );
1745    }
1746
1747    #[test]
1748    fn nested_with_panics_instead_of_deadlocking() {
1749        let result = within_deadline(|| {
1750            let live = live(vec![ProgressColumn::Description]);
1751            let nested = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
1752                live.with(|_| live.add_task("nested", None, 0.0))
1753            }));
1754            // The display is still usable after the refused re-entry.
1755            let task = live.add_task("after", None, 0.0);
1756            live.stop();
1757            (nested.is_err(), task)
1758        });
1759        assert!(result.0, "a nested `with` must be refused, not deadlock");
1760    }
1761
1762    #[test]
1763    fn track_counts_each_item_after_its_loop_body() {
1764        let live = live(vec![ProgressColumn::Description, ProgressColumn::MofN]);
1765        let mut seen = Vec::new();
1766        let tracked = live.track(vec!['a', 'b', 'c'], None, "letters");
1767        let task = tracked.task();
1768        for item in tracked {
1769            // Upstream advances after the loop body, so the item being
1770            // processed is not yet counted.
1771            let completed = live.with(|progress| progress.task(task).unwrap().completed());
1772            seen.push((item, completed));
1773        }
1774        assert_eq!(seen, vec![('a', 0.0), ('b', 1.0), ('c', 2.0)]);
1775        let (progress, bytes) = live.stop();
1776        let task = progress.task(task).unwrap();
1777        assert_eq!((task.total(), task.completed()), (Some(3.0), 3.0));
1778        let output = String::from_utf8(bytes).unwrap();
1779        assert!(
1780            output.ends_with("letters \x1b[32m3/3\x1b[0m\n\x1b[?25h"),
1781            "{output:?}"
1782        );
1783    }
1784
1785    #[test]
1786    fn track_leaves_an_iterator_of_unknown_length_indeterminate() {
1787        let live = live(vec![ProgressColumn::MofN]);
1788        let task = {
1789            let mut tracked = live.track((0..10).filter(|n| n % 3 == 0), None, "");
1790            let task = tracked.task();
1791            assert_eq!(tracked.by_ref().count(), 4);
1792            task
1793        };
1794        let with_total = {
1795            let mut tracked = live.track(0..2, Some(5.0), "");
1796            tracked.by_ref().for_each(drop);
1797            tracked.task()
1798        };
1799        let (progress, _) = live.stop();
1800        assert_eq!(progress.task(task).unwrap().total(), None);
1801        assert_eq!(progress.task(task).unwrap().completed(), 4.0);
1802        assert_eq!(progress.task(with_total).unwrap().total(), Some(5.0));
1803    }
1804
1805    fn quiet_console() -> Console {
1806        Console::builder().force_terminal(false).width(40).build()
1807    }
1808
1809    #[test]
1810    fn wrap_read_advances_by_the_bytes_read() {
1811        use std::io::Read;
1812        let live = Progress::new()
1813            .disable(true)
1814            .start(quiet_console(), Vec::new(), 1.0);
1815        let mut reader = live
1816            .wrap_read(&b"hello world"[..], Some(11), None, "Reading...")
1817            .expect("total given");
1818        let task = reader.task();
1819        let mut buf = [0u8; 4];
1820        reader.read_exact(&mut buf).unwrap();
1821        assert_eq!(live.with(|p| p.task(task).unwrap().completed()), 4.0);
1822        let mut rest = Vec::new();
1823        reader.read_to_end(&mut rest).unwrap();
1824        assert!(live.with(|p| p.task(task).unwrap().finished()));
1825        let (_, out) = live.stop();
1826        assert!(out.is_empty(), "a disabled display writes nothing");
1827    }
1828
1829    #[test]
1830    fn wrap_read_needs_a_total() {
1831        let live = Progress::new()
1832            .disable(true)
1833            .start(quiet_console(), Vec::new(), 1.0);
1834        let err = live
1835            .wrap_read(&b""[..], None, None, "x")
1836            .err()
1837            .expect("no total");
1838        assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput);
1839        let task = live.add_task("sized", 5.0, 0.0);
1840        assert!(live.wrap_read(&b""[..], None, Some(task), "x").is_ok());
1841    }
1842
1843    #[test]
1844    fn open_takes_the_file_size_as_total() {
1845        use std::io::Read;
1846        let path = std::env::temp_dir().join(format!("rs-rich-open-{}", std::process::id()));
1847        std::fs::write(&path, b"0123456789").unwrap();
1848        let live = Progress::new()
1849            .disable(true)
1850            .start(quiet_console(), Vec::new(), 1.0);
1851        let mut reader = live.open(&path, None, None, "Reading...").unwrap();
1852        let task = reader.task();
1853        assert_eq!(live.with(|p| p.task(task).unwrap().total()), Some(10.0));
1854        std::io::copy(&mut reader, &mut std::io::sink()).unwrap();
1855        assert_eq!(live.with(|p| p.task(task).unwrap().completed()), 10.0);
1856        let _ = reader.read(&mut [0u8; 1]);
1857        drop(reader);
1858        std::fs::remove_file(path).unwrap();
1859    }
1860}