hf_fetch_model/
progress.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Progress reporting for model downloads.
4//!
5//! [`ProgressEvent`] carries per-file and overall download status.
6//! When the `indicatif` feature is enabled, `IndicatifProgress`
7//! provides multi-progress bars out of the box.
8
9/// A progress event emitted during download.
10///
11/// Passed to the `on_progress` callback on [`crate::FetchConfig`].
12#[derive(Debug, Clone)]
13pub struct ProgressEvent {
14    /// The filename currently being downloaded.
15    pub filename: String,
16    /// Bytes downloaded so far for this file.
17    pub bytes_downloaded: u64,
18    /// Total size of this file in bytes (0 if unknown).
19    pub bytes_total: u64,
20    /// Download percentage for this file (0.0–100.0).
21    pub percent: f64,
22    /// Number of files still remaining (after this one).
23    pub files_remaining: usize,
24}
25
26/// Creates a [`ProgressEvent`] for a completed file.
27#[must_use]
28pub(crate) fn completed_event(filename: &str, size: u64, files_remaining: usize) -> ProgressEvent {
29    ProgressEvent {
30        filename: filename.to_owned(),
31        bytes_downloaded: size,
32        bytes_total: size,
33        percent: 100.0,
34        files_remaining,
35    }
36}
37
38/// Creates a [`ProgressEvent`] for an in-progress file (streaming update).
39#[must_use]
40pub(crate) fn streaming_event(
41    filename: &str,
42    bytes_downloaded: u64,
43    bytes_total: u64,
44    files_remaining: usize,
45) -> ProgressEvent {
46    #[allow(clippy::cast_precision_loss, clippy::as_conversions)]
47    // EXPLICIT: f64 cast for percentage; precision loss negligible for display
48    let percent = if bytes_total > 0 {
49        (bytes_downloaded as f64 / bytes_total as f64) * 100.0
50    } else {
51        0.0
52    };
53    ProgressEvent {
54        filename: filename.to_owned(),
55        bytes_downloaded,
56        bytes_total,
57        percent,
58        files_remaining,
59    }
60}
61
62/// Multi-progress bar display using `indicatif`.
63///
64/// Available only when the `indicatif` feature is enabled.
65///
66/// # Example
67///
68/// ```rust,no_run
69/// # fn example() -> Result<(), hf_fetch_model::FetchError> {
70/// use hf_fetch_model::FetchConfig;
71/// # #[cfg(feature = "indicatif")]
72/// use hf_fetch_model::progress::IndicatifProgress;
73///
74/// # #[cfg(feature = "indicatif")]
75/// let progress = IndicatifProgress::new();
76/// let config = FetchConfig::builder()
77///     # ;
78///     # #[cfg(feature = "indicatif")]
79///     # let config = FetchConfig::builder()
80///     .on_progress(move |e| progress.handle(e))
81///     .build()?;
82/// # Ok(())
83/// # }
84/// ```
85#[cfg(feature = "indicatif")]
86pub struct IndicatifProgress {
87    multi: indicatif::MultiProgress,
88    overall: indicatif::ProgressBar,
89    finished: std::sync::atomic::AtomicBool,
90}
91
92#[cfg(feature = "indicatif")]
93impl IndicatifProgress {
94    /// Creates a new multi-progress bar display.
95    ///
96    /// Call [`IndicatifProgress::set_total_files`] once the file count is known.
97    #[must_use]
98    pub fn new() -> Self {
99        let multi = indicatif::MultiProgress::new();
100        let overall = multi.add(indicatif::ProgressBar::new(0));
101        overall.set_style(
102            indicatif::ProgressStyle::default_bar()
103                .template("{msg} [{bar:40.cyan/blue}] {pos}/{len} files")
104                .ok()
105                .unwrap_or_else(indicatif::ProgressStyle::default_bar)
106                .progress_chars("=> "),
107        );
108        overall.set_message("Overall");
109        Self {
110            multi,
111            overall,
112            finished: std::sync::atomic::AtomicBool::new(false),
113        }
114    }
115
116    /// Returns a reference to the underlying [`indicatif::MultiProgress`].
117    ///
118    /// Useful for adding custom progress bars alongside the built-in ones.
119    #[must_use]
120    pub fn multi(&self) -> &indicatif::MultiProgress {
121        &self.multi
122    }
123
124    /// Sets the total number of files to download.
125    pub fn set_total_files(&self, total: u64) {
126        self.overall.set_length(total);
127    }
128
129    /// Handles a [`ProgressEvent`], updating progress bars.
130    pub fn handle(&self, event: &ProgressEvent) {
131        if event.percent >= 100.0 {
132            // Derive total: completed so far + this file + remaining
133            // EXPLICIT: try_from for usize → u64 (infallible on 64-bit, safe fallback otherwise)
134            let remaining = u64::try_from(event.files_remaining).unwrap_or(u64::MAX);
135            let total = self.overall.position() + 1 + remaining;
136            self.overall.set_length(total);
137            self.overall.inc(1);
138        }
139    }
140
141    /// Finishes the progress bar, ensuring the final state is rendered.
142    ///
143    /// Called automatically on drop, but can be called explicitly for
144    /// immediate visual feedback.
145    pub fn finish(&self) {
146        if !self
147            .finished
148            .swap(true, std::sync::atomic::Ordering::Relaxed)
149        {
150            self.overall.finish();
151        }
152    }
153}
154
155#[cfg(feature = "indicatif")]
156impl Drop for IndicatifProgress {
157    fn drop(&mut self) {
158        self.finish();
159    }
160}
161
162#[cfg(feature = "indicatif")]
163impl Default for IndicatifProgress {
164    fn default() -> Self {
165        Self::new()
166    }
167}
hf_fetch_model/progress.rs

hf_fetch_model/
progress.rs
