Skip to main content

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/// Multi-progress bar display using `indicatif`.
39///
40/// Available only when the `indicatif` feature is enabled.
41///
42/// # Example
43///
44/// ```rust,no_run
45/// # fn example() -> Result<(), hf_fetch_model::FetchError> {
46/// use hf_fetch_model::FetchConfig;
47/// # #[cfg(feature = "indicatif")]
48/// use hf_fetch_model::progress::IndicatifProgress;
49///
50/// # #[cfg(feature = "indicatif")]
51/// let progress = IndicatifProgress::new();
52/// let config = FetchConfig::builder()
53///     # ;
54///     # #[cfg(feature = "indicatif")]
55///     # let config = FetchConfig::builder()
56///     .on_progress(move |e| progress.handle(e))
57///     .build()?;
58/// # Ok(())
59/// # }
60/// ```
61#[cfg(feature = "indicatif")]
62pub struct IndicatifProgress {
63    multi: indicatif::MultiProgress,
64    overall: indicatif::ProgressBar,
65}
66
67#[cfg(feature = "indicatif")]
68impl IndicatifProgress {
69    /// Creates a new multi-progress bar display.
70    ///
71    /// Call [`IndicatifProgress::set_total_files`] once the file count is known.
72    #[must_use]
73    pub fn new() -> Self {
74        let multi = indicatif::MultiProgress::new();
75        let overall = multi.add(indicatif::ProgressBar::new(0));
76        overall.set_style(
77            indicatif::ProgressStyle::default_bar()
78                .template("{msg} [{bar:40.cyan/blue}] {pos}/{len} files")
79                .ok()
80                .unwrap_or_else(indicatif::ProgressStyle::default_bar)
81                .progress_chars("=> "),
82        );
83        overall.set_message("Overall");
84        Self { multi, overall }
85    }
86
87    /// Returns a reference to the underlying [`indicatif::MultiProgress`].
88    ///
89    /// Useful for adding custom progress bars alongside the built-in ones.
90    #[must_use]
91    pub fn multi(&self) -> &indicatif::MultiProgress {
92        &self.multi
93    }
94
95    /// Sets the total number of files to download.
96    pub fn set_total_files(&self, total: u64) {
97        self.overall.set_length(total);
98    }
99
100    /// Handles a [`ProgressEvent`], updating progress bars.
101    pub fn handle(&self, event: &ProgressEvent) {
102        if event.percent >= 100.0 {
103            // Derive total: completed so far + this file + remaining
104            // EXPLICIT: try_from for usize → u64 (infallible on 64-bit, safe fallback otherwise)
105            let remaining = u64::try_from(event.files_remaining).unwrap_or(u64::MAX);
106            let total = self.overall.position() + 1 + remaining;
107            self.overall.set_length(total);
108            self.overall.inc(1);
109        }
110    }
111}
112
113#[cfg(feature = "indicatif")]
114impl Default for IndicatifProgress {
115    fn default() -> Self {
116        Self::new()
117    }
118}