From c8c3290f3250482e6dd8f5f60730f213158e35d8 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Mon, 13 Jun 2022 08:44:16 +0000 Subject: [PATCH] Update --- src/downloader/downloader.rs | 119 ++++++++++++++++++++++++++++++++++- src/downloader/error.rs | 2 + src/downloader/local_file.rs | 3 + src/downloader/mod.rs | 1 + src/downloader/tasks.rs | 18 +++++- src/utils.rs | 1 + 6 files changed, 140 insertions(+), 4 deletions(-) diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index b44b571..614f825 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -10,10 +10,16 @@ use crate::ext::io::ClearFile; use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; use crate::ext::try_err::TryErr; +use crate::gettext; use crate::utils::ask_need_overwrite; +use crate::utils::get_file_name_from_url; use crate::webclient::WebClient; use crate::webclient::ToHeaders; +use indicatif::ProgressBar; +use indicatif::ProgressStyle; +use indicatif::MultiProgress; use reqwest::IntoUrl; +use std::borrow::Cow; use std::collections::HashMap; #[cfg(test)] use std::fs::create_dir; @@ -21,6 +27,7 @@ use std::fs::remove_file; use std::io::Seek; use std::io::SeekFrom; use std::io::Write; +use std::ops::Deref; use std::ops::DerefMut; use std::path::Path; use std::sync::Arc; @@ -48,6 +55,10 @@ pub struct DownloaderInternal { pub tasks: RwLock>>>, /// Whether to enable mulitple thread mode multi: AtomicBool, + /// Whether to enable the progess bar. + progress_bar: AtomicBool, + /// The progress bar. + progress: RwLock>, } impl DownloaderInternal { @@ -114,6 +125,8 @@ impl DownloaderInternal { status: RwLock::new(DownloaderStatus::Created), tasks: RwLock::new(Vec::new()), multi: AtomicBool::new(false), + progress_bar: AtomicBool::new(false), + progress: RwLock::new(None), })) } } @@ -134,11 +147,64 @@ impl DownloaderInternal { Ok(()) } + /// Disable the progress bar + pub fn disable_progress_bar(&self) { + self.progress_bar.qstore(false); + self.progress.get_mut().take(); + } + + /// Enable the progress bar + /// * `style` - The style of the progress bar + /// * `mults` - The instance of [MultiProgress] if multiple progress bars are needed. + pub fn enable_progress_bar(&self, style: ProgressStyle, mults: Option<&MultiProgress>) { + let mut bar = ProgressBar::new(0).with_style(style); + match mults { + Some(bars) => { + bar = bars.add(bar); + } + None => { } + } + self.progress_bar.qstore(true); + self.progress.get_mut().replace(bar); + } + + #[inline] + /// Returns true if the progress bar is enabled. + pub fn enabled_progress_bar(&self) -> bool { + self.progress_bar.qload() + } + + #[inline] + /// Finishes the progress bar and sets a message + pub fn finish_progress_bar_with_message(&self, msg: impl Into>) { + match self.progress.get_ref().deref() { + Some(p) => { p.finish_with_message(msg) } + None => { } + } + } + + #[inline] + /// Get the file name from url. + /// If not available, use `(Unknown)` + pub fn get_file_name(&self) -> String { + get_file_name_from_url(self.url.deref().clone()).unwrap_or(String::from(gettext("(Unknown)"))) + } + + #[inline] /// Return the status of the downloader. pub fn get_status(&self) -> DownloaderStatus { self.status.get_ref().clone() } + #[inline] + /// Advances the position of the progress bar by `delta` + pub fn inc_progress_bar(&self, delta: u64) { + match self.progress.get_ref().deref() { + Some(p) => { p.inc(delta) } + None => { } + } + } + #[inline] /// Returns true if the downloader is created just now. pub fn is_created(&self) -> bool { @@ -151,6 +217,12 @@ impl DownloaderInternal { *self.status.get_ref() == DownloaderStatus::Downloading } + #[inline] + /// Returns true if the downloader is downloaded complete. + pub fn is_downloaded(&self) -> bool { + *self.status.get_ref() == DownloaderStatus::Downloaded + } + #[inline] /// Returns true if is multiple thread mode. pub fn is_multi_threads(&self) -> bool { @@ -182,6 +254,33 @@ impl DownloaderInternal { self.status.replace_with2(DownloaderStatus::Downloaded) } + #[inline] + /// Sets the length of the progress bar + pub fn set_progress_bar_length(&self, length: u64) { + match self.progress.get_ref().deref() { + Some(p) => { p.set_length(length) } + None => { } + } + } + + #[inline] + /// Sets the position of the progress bar + pub fn set_progress_bar_position(&self, pos: u64) { + match self.progress.get_ref().deref() { + Some(p) => { p.set_position(pos) } + None => { } + } + } + + #[inline] + /// Sets the current message of the progress bar + pub fn set_progress_bar_message(&self, msg: impl Into>) { + match self.progress.get_ref().deref() { + Some(p) => { p.set_message(msg) } + None => { } + } + } + /// Write datas to the file. /// * `data` - Data pub fn write(&self, data: &[u8]) -> Result<(), DownloaderError> { @@ -241,19 +340,31 @@ impl Downloader { self.downloader.get_status() } + /// Wait the downloader. pub async fn join(&self) -> Result<(), DownloaderError> { match self.task.get_mut().deref_mut() { Some(v) => { - let re = v.await; - re.unwrap() + v.await? } None => { Ok(()) } } } + /// Disable progress bar + pub fn disable_progress_bar(&self) { + self.downloader.disable_progress_bar() + } + + /// Enable the progress bar + /// * `style` - The style of the progress bar + /// * `mults` - The instance of [MultiProgress] if multiple progress bars are needed. + pub fn enable_progress_bar(&self, style: ProgressStyle, mults: Option<&MultiProgress>) { + self.downloader.enable_progress_bar(style, mults) + } define_downloader_fn!(is_created, bool, "Returns true if the downloader is created just now."); define_downloader_fn!(is_downloading, bool, "Returns true if the downloader is downloading now."); define_downloader_fn!(is_multi_threads, bool, "Returns true if is multiple thread mode."); + define_downloader_fn!(is_downloaded, bool, "Returns true if the downloader is downloaded complete."); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -269,8 +380,10 @@ async fn test_downloader() { match downloader { DownloaderResult::Ok(v) => { assert_eq!(v.is_created(), true); + v.disable_progress_bar(); v.download(); - assert!(v.join().await.is_ok()); + v.join().await.unwrap(); + assert_eq!(v.is_downloaded(), true); } DownloaderResult::Canceled => { panic!("This should not happened.") } } diff --git a/src/downloader/error.rs b/src/downloader/error.rs index 76fe3cf..bd7bac5 100644 --- a/src/downloader/error.rs +++ b/src/downloader/error.rs @@ -1,6 +1,7 @@ use crate::downloader::pd_file::PdFileError; use http::status::StatusCode; use tokio::time::error::Elapsed; +use tokio::task::JoinError; #[derive(Debug, derive_more::From)] pub enum DownloaderError { @@ -10,6 +11,7 @@ pub enum DownloaderError { String(String), ErrorStatusCode(StatusCode), Timeout(Elapsed), + JoinError(JoinError), } impl From<&str> for DownloaderError { diff --git a/src/downloader/local_file.rs b/src/downloader/local_file.rs index 8d632a4..76731a9 100644 --- a/src/downloader/local_file.rs +++ b/src/downloader/local_file.rs @@ -7,8 +7,11 @@ use std::ops::Deref; use std::path::Path; use std::path::PathBuf; +/// A wrapper for [File], add support for clear file content. pub struct LocalFile { + /// The file. file: Option, + /// The path of the file. path: PathBuf, } diff --git a/src/downloader/mod.rs b/src/downloader/mod.rs index 61f1cbf..b047961 100644 --- a/src/downloader/mod.rs +++ b/src/downloader/mod.rs @@ -12,3 +12,4 @@ pub mod pd_file; pub mod tasks; pub use downloader::Downloader; pub use error::DownloaderError; +pub use local_file::LocalFile; diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs index 9140a30..843cb5b 100644 --- a/src/downloader/tasks.rs +++ b/src/downloader/tasks.rs @@ -90,10 +90,16 @@ pub async fn create_download_tasks_simple { d.pd.set_file_size(len)?; + if d.enabled_progress_bar() { + d.set_progress_bar_length(len); + } } None => {} } } + if d.enabled_progress_bar() { + d.set_progress_bar_message(gettext("Downloading \"\".").replace("", d.get_file_name().as_str())); + } handle_download(d, result).await } @@ -111,13 +117,20 @@ pub async fn handle_download(d: Arc { if !is_multi { - d.pd.inc(data.len() as u64)?; + let len = data.len() as u64; + d.pd.inc(len)?; + if d.enabled_progress_bar() { + d.inc_progress_bar(len); + } } d.write(&data)?; } Err(e) => { if !is_multi { d.pd.clear()?; + if d.enabled_progress_bar() { + d.set_progress_bar_position(0); + } } return Err(DownloaderError::from(e)); } @@ -134,6 +147,9 @@ pub async fn handle_download(d: Arc { if !is_multi { d.pd.clear()?; + if d.enabled_progress_bar() { + d.set_progress_bar_position(0); + } } return Err(DownloaderError::from(e)); } diff --git a/src/utils.rs b/src/utils.rs index 2635bf1..24962cc 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -45,6 +45,7 @@ pub fn ask_need_overwrite(path: &str) -> bool { } } +/// Get file name from url. pub fn get_file_name_from_url(url: U) -> Option { let u = url.into_url(); if u.is_err() {