From 401341047f206242a3efe6ac1f58a0baf728f1ad Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 17 Jul 2022 11:58:20 +0000 Subject: [PATCH] Add support to fanbox.creator url embed type --- src/download.rs | 74 +++++++++++++++++++++++++++++++++++++++++++-- src/task_manager.rs | 34 +++++++++++++++++++++ 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/src/download.rs b/src/download.rs index 0e7310d..a18aee7 100644 --- a/src/download.rs +++ b/src/download.rs @@ -16,6 +16,7 @@ use crate::error::PixivDownloaderError; use crate::ext::any::AsAny; use crate::ext::try_err::TryErr; use crate::fanbox::article::block::FanboxArticleBlock; +use crate::fanbox::article::url_embed::FanboxArticleUrlEmbed; use crate::fanbox::check::CheckUnknown; use crate::fanbox::creator::FanboxCreator; use crate::fanbox::creator::FanboxProfileItem; @@ -109,6 +110,7 @@ impl Main { Arc::clone(&fc), id.to_owned(), None, + None, )) .await; if !download_multiple_posts { @@ -584,6 +586,8 @@ pub async fn download_fanbox_post( .save(&json_file) .try_err(gettext("Failed to save post data to file."))?; let tasks = TaskManager::default(); + let ptasks = TaskManager::new_post(); + let mut re = Ok(()); let download_multiple_files = helper.download_multiple_files(); match post { FanboxPost::Article(article) => { @@ -595,6 +599,9 @@ pub async fn download_fanbox_post( let blocks = body .blocks() .try_err(gettext("Failed to get blocks from article."))?; + let url_embed_map = body + .url_embed_map() + .try_err(gettext("Failed to get embed url map from article."))?; let mut np = 0; let mut datas = data.clone(); #[cfg(feature = "exif")] @@ -626,6 +633,43 @@ pub async fn download_fanbox_post( } np += 1; } + FanboxArticleBlock::UrlEmbed(u) => { + let embed_url = url_embed_map + .get_url_embed( + u.url_embed_id() + .try_err(gettext("Failed to get embed url id from block"))?, + ) + .try_err(gettext("Failed to get embed url from url embed map."))?; + match embed_url { + FanboxArticleUrlEmbed::FanboxCreator(creator) => { + let profile = creator + .profile() + .try_err(gettext("Failed to get creator's profile."))?; + let id = profile + .creator_id() + .try_err("Failed to get creator's id.")?; + match ptasks + .add_task_else_run_local(download_fanbox_creator_info( + Arc::clone(&fc), + id.to_owned(), + Some(profile), + None, + )) + .await + { + Some(r) => { + concat_pixiv_downloader_error!(re, r); + } + None => { + if !download_multiple_files { + ptasks.join().await; + } + } + } + } + _ => {} + } + } _ => {} } } @@ -692,7 +736,6 @@ pub async fn download_fanbox_post( } } tasks.join().await; - let mut re = Ok(()); let tasks = tasks.take_finished_tasks(); for mut task in tasks { let task = task.as_any_mut(); @@ -705,6 +748,19 @@ pub async fn download_fanbox_post( concat_pixiv_downloader_error!(re, r); } } + ptasks.join().await; + let ptasks = ptasks.take_finished_tasks(); + for mut task in ptasks { + let task = task.as_any_mut(); + if let Some(task) = task.downcast_mut::>>() { + let r = task.await; + let r = match r { + Ok(r) => r, + Err(e) => Err(PixivDownloaderError::from(e)), + }; + concat_pixiv_downloader_error!(re, r); + } + } re } @@ -712,6 +768,7 @@ pub async fn download_fanbox_creator_info( fc: Arc, id: String, data: Option, + base: Option, ) -> Result<(), PixivDownloaderError> { let data = match data { Some(data) => { @@ -751,9 +808,22 @@ pub async fn download_fanbox_creator_info( } let mut fdata = FanboxData::new(PixivID::FanboxCreator(id.clone()), &*data) .try_err("Failed to create data file.")?; - let base = Arc::new(PathBuf::from(".").join(&id)); + let base = match base { + Some(base) => Arc::new(base), + None => Arc::new(PathBuf::from(".").join(&id)), + }; let json_file = base.join("creator.json"); let data_file = JSONDataFile::from(&fdata); + if !base.exists() { + match create_dir_all(&*base) { + Ok(_) => {} + Err(e) => { + if !base.exists() { + return Err(PixivDownloaderError::from(e)); + } + } + } + } data_file .save(&json_file) .try_err(gettext("Failed to save post data to file."))?; diff --git a/src/task_manager.rs b/src/task_manager.rs index 3ab0b26..d3ef50d 100644 --- a/src/task_manager.rs +++ b/src/task_manager.rs @@ -147,6 +147,40 @@ impl TaskManager { } } + /// Try add a new task, if queue is full, run future on local thread + pub async fn add_task_else_run_local(&self, future: F) -> Option + where + F: Future + Send + 'static, + F::Output: Send + 'static, + JoinHandle: IsFinishedAny + Send + Sync + 'static, + { + let total_count = self.max_count.get_max_count(); + { + let mut count = self.task_count.lock().await; + let tasks = self.tasks.replace_with2(Vec::new()); + let mut new_tasks = Vec::new(); + let mut new_count = *count; + for i in tasks { + if i.is_finished() { + self.finished_tasks.get_mut().push(i); + new_count -= 1; + } else { + new_tasks.push(i); + } + } + self.tasks.replace_with2(new_tasks); + count.replace_with(new_count); + if *count < total_count { + self.tasks + .get_mut() + .push(Box::new(tokio::task::spawn(future))); + count.replace_with(*count + 1); + return None; + } + } + Some(future.await) + } + /// Wait all tasks finished. pub async fn join(&self) { loop {