use std::{ error, fmt, fs::File, io::{self, Read, Write}, }; use bytes::Bytes; use futures_util::StreamExt; use indicatif::{ProgressBar, ProgressState, ProgressStyle}; use reqwest::{self, Client}; use std::cmp; use crate::get_mut_ref; pub enum DualWriter { File(File), Buffer(Vec), } impl DualWriter { pub fn write(&mut self, bytes: Bytes) -> Result<(), std::io::Error> { match self { Self::Buffer(buffer) => { bytes.into_iter().for_each(|byte| buffer.push(byte)); } Self::File(file) => { file.write(&bytes)?; } } Ok(()) } } impl TryFrom> for DualWriter { type Error = io::Error; fn try_from(file_name: Option<&str>) -> Result { Ok(if let Some(file_name) = file_name { Self::File(File::create(&file_name)?) } else { Self::Buffer(Vec::::new()) }) } } impl TryInto for DualWriter { type Error = String; fn try_into(self) -> Result { Ok(match self { Self::Buffer(buffer) => { String::from_utf8(buffer).or(Err("Failed to decode buffer".to_owned()))? } Self::File(file) => { let mut buf = String::new(); // Well this is safe since I consume the file anyways let file = unsafe { get_mut_ref(&file) }; file.read_to_string(&mut buf) .or(Err("Failed to read file".to_owned()))?; buf } }) } } pub async fn download_with_progress( link: &str, file_name: Option<&str>, ) -> Result> { let mut dual_writer: DualWriter = file_name.try_into()?; let client = Client::builder().gzip(true).deflate(true).build()?; let resp = client.get(link).send().await?; let content_length = if let Some(got_content_length) = resp.content_length() { got_content_length } else { panic!("Could not retrive content length from server. {:?}", &resp); }; let progress_bar = ProgressBar::new(content_length); progress_bar.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})").unwrap() .with_key("eta", |state: &ProgressState, w: &mut dyn fmt::Write| write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()) .progress_chars("#>-")); let mut downloaded: u64 = 0; let mut stream = resp.bytes_stream(); while let Some(item) = stream.next().await { let bytes = item.unwrap(); downloaded = cmp::min(downloaded + (bytes.len() as u64), content_length); dual_writer.write(bytes)?; progress_bar.set_position(downloaded); } Ok(dual_writer) }