Skip to content

Commit

Permalink
Merge pull request #6098 from horriblename/feat-download-progress
Browse files Browse the repository at this point in the history
Add simple download progress tracker
  • Loading branch information
bhansconnect authored Dec 2, 2023
2 parents d21df40 + edaa600 commit d982764
Showing 1 changed file with 50 additions and 2 deletions.
52 changes: 50 additions & 2 deletions crates/packaging/src/https.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::{
io::{self, Read},
io::{self, Read, Write},
path::Path,
};

Expand Down Expand Up @@ -275,9 +275,12 @@ pub fn download_and_hash(
Encoding::new(content_encoding, url)?
};

let content_length = resp.content_length().map(|n| n as usize);

// Use .take to prevent a malicious server from sending back bytes
// until system resources are exhausted!
decompress_into(dest_dir, encoding, resp.take(max_download_bytes))
let resp = ProgressReporter::new(resp.take(max_download_bytes), content_length);
decompress_into(dest_dir, encoding, resp)
}

/// The content encodings we support
Expand Down Expand Up @@ -405,3 +408,48 @@ impl<R: Read> Read for HashReader<R> {
Ok(bytes_read)
}
}

/// Prints download progress to stdout
struct ProgressReporter<R: Read> {
read: usize,
total: Option<usize>,
reader: R,
}

impl<R: Read> ProgressReporter<R> {
fn new(reader: R, total: Option<usize>) -> Self {
ProgressReporter {
read: 0,
total,
reader,
}
}
}

impl<R: Read> Read for ProgressReporter<R> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let size = self.reader.read(buf)?;

self.read += size;

if let Some(total) = self.total {
eprint!(
"\u{001b}[2K\u{001b}[G[{:.1} / {:.1} MB]",
self.read as f32 / 1_000_000.0,
total as f32 / 1_000_000.0,
);
} else {
eprint!(
"\u{001b}[2K\u{001b}[G[{:.1} MB]",
self.read as f32 / 1_000_000.0,
);
}
std::io::stderr().flush()?;

if self.total.is_some_and(|total| self.read >= total) {
eprintln!();
}

Ok(size)
}
}

0 comments on commit d982764

Please sign in to comment.