From ecf883e6e7231bcc4d3ef52194c4e10738d60ed2 Mon Sep 17 00:00:00 2001 From: trivernis Date: Sat, 13 Aug 2022 13:43:27 +0200 Subject: [PATCH] Refactor entry point Signed-off-by: trivernis --- src/main.rs | 162 +++++++-------------------- src/operations/find_and_send_tags.rs | 65 +++++++++++ src/operations/find_and_send_urls.rs | 21 ++++ src/operations/mod.rs | 2 + src/utils/mod.rs | 1 + src/{search.rs => utils/pixiv.rs} | 22 +++- 6 files changed, 146 insertions(+), 127 deletions(-) create mode 100644 src/operations/find_and_send_tags.rs create mode 100644 src/operations/find_and_send_urls.rs create mode 100644 src/operations/mod.rs create mode 100644 src/utils/mod.rs rename src/{search.rs => utils/pixiv.rs} (63%) diff --git a/src/main.rs b/src/main.rs index 7063434..7fa3c31 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,17 +1,18 @@ mod args; mod error; -pub mod search; +mod operations; +pub mod utils; use crate::error::Result; -use crate::search::get_urls; +use crate::operations::find_and_send_tags::find_and_send_tags; +use crate::operations::find_and_send_urls::find_and_send_urls; use args::*; use clap::Parser; -use hydrus_api::wrapper::hydrus_file::HydrusFile; use hydrus_api::wrapper::service::ServiceName; use hydrus_api::wrapper::tag::Tag; use hydrus_api::{Client, Hydrus}; use pixiv_rs::PixivClient; -use rustnao::{Handler, HandlerBuilder, Sauce}; +use rustnao::{Handler, HandlerBuilder}; use std::str::FromStr; use tempdir::TempDir; use tokio::time::{Duration, Instant}; @@ -23,53 +24,13 @@ async fn main() { init_logger(); let args: Args = Args::parse(); tracing::debug!("args: {args:?}"); - let opt = match &args.subcommand { - Command::FindAndSendUrl(opt) => opt.clone(), - Command::FindAndSendTags(opt) => opt.clone(), - }; - - let handler = HandlerBuilder::new() - .api_key(&opt.saucenao_key) - .min_similarity(80) - .db(Handler::PIXIV) - .build(); - let hydrus = Hydrus::new(Client::new(&args.hydrus_url, &args.hydrus_key)); - let pixiv = PixivClient::new(); - - let tags = opt.tags.into_iter().map(Tag::from).collect(); - let service = ServiceName(opt.tag_service); - - let files = hydrus.search().add_tags(tags).run().await.unwrap(); - tracing::info!("Found {} files", files.len()); - let tmpdir = TempDir::new("hydrus-files").unwrap(); - - let sleep_duration = Duration::from_secs(6); - - for mut file in files { - let start = Instant::now(); - match &args.subcommand { - Command::FindAndSendUrl(_) => { - let _ = search_and_send_urls(&hydrus, &handler, &tmpdir, &mut file).await; - } - Command::FindAndSendTags(_) => { - let _ = tag_file( - opt.finish_tag.as_ref(), - &handler, - &pixiv, - &service, - &tmpdir, - &mut file, - ) - .await; - } - } - let elapsed = start.elapsed(); - if elapsed.as_secs() < 8 { - tokio::time::sleep(sleep_duration - elapsed).await; // rate limit of 6# / 30s - } + match args.subcommand { + Command::FindAndSendUrl(opt) => send_tags_or_urls(opt, hydrus, true).await, + Command::FindAndSendTags(opt) => send_tags_or_urls(opt, hydrus, false).await, } + .expect("Failed to send tags or urls"); } fn init_logger() { @@ -86,87 +47,44 @@ fn init_logger() { .init(); } -async fn tag_file( - finish_tag: Option<&String>, - handler: &Handler, - pixiv: &PixivClient, - service: &ServiceName, - tmpdir: &TempDir, - mut file: &mut HydrusFile, -) -> Result<()> { - if let Err(e) = search_and_assign_tags(&handler, &pixiv, &service, &tmpdir, &mut file).await { - let hash = file.hash().await.unwrap(); - tracing::error!("Failed to search tags to file {}: {:?}", hash, e); - } else if let Some(finish_tag) = finish_tag { - file.add_tags(service.clone().into(), vec![finish_tag.into()]) - .await - .unwrap(); - } - - Ok(()) -} - -async fn search_and_send_urls( - hydrus: &Hydrus, - handler: &Handler, - tmpdir: &TempDir, - file: &mut HydrusFile, -) -> Result<()> { - let sauces = get_sauces_for_file(&handler, tmpdir, file).await?; - let urls = get_urls(&sauces); - for url in urls { - hydrus.import().url(url).run().await?; - } - - Ok(()) -} - -async fn search_and_assign_tags( - handler: &Handler, - pixiv: &PixivClient, - service: &ServiceName, - tmpdir: &TempDir, - mut file: &mut HydrusFile, -) -> Result<()> { - tracing::debug!("Getting tags for hydrus file {:?}", file.id); - let sauces = get_sauces_for_file(&handler, tmpdir, file).await?; +async fn send_tags_or_urls(opt: Options, hydrus: Hydrus, send_urls: bool) -> Result<()> { + let pixiv = PixivClient::new(); - assign_pixiv_tags_and_url(&pixiv, service, &mut file, &sauces).await -} + let handler = HandlerBuilder::new() + .api_key(&opt.saucenao_key) + .min_similarity(80) + .db(Handler::PIXIV) + .build(); -async fn get_sauces_for_file( - handler: &Handler, - tmpdir: &TempDir, - mut file: &mut HydrusFile, -) -> Result> { - tracing::debug!("Creating tmp file for hydrus file {:?}", file.id); - let path = search::create_tmp_sauce_file(&tmpdir, &mut file).await?; - tracing::debug!("Getting sauce for hydrus file {:?}", file.id); + let tags = opt.tags.into_iter().map(Tag::from).collect(); + let service = ServiceName(opt.tag_service); - let sauce = handler.get_sauce(path.to_str().unwrap(), None, None)?; - tracing::debug!("Getting tags for hydrus file {:?}", file.id); - Ok(sauce) -} + let files = hydrus.search().add_tags(tags).run().await.unwrap(); + tracing::info!("Found {} files", files.len()); + let tmpdir = TempDir::new("hydrus-files").unwrap(); -async fn assign_pixiv_tags_and_url( - pixiv: &&PixivClient, - service: &ServiceName, - file: &mut &mut HydrusFile, - sauce: &Vec, -) -> Result<()> { - let hash = file.hash().await?; - if let Some(url) = search::get_pixiv_url(&sauce) { - let tags = search::get_tags_for_sauce(&pixiv, url).await?; + let sleep_duration = Duration::from_secs(6); - if tags.len() > 0 { - tracing::info!("Found {} tags for file {:?}", tags.len(), hash); - file.add_tags(service.clone().into(), tags).await?; + for mut file in files { + let start = Instant::now(); + if send_urls { + let _ = find_and_send_urls(&hydrus, &handler, &tmpdir, &mut file).await; } else { - tracing::info!("No tags for file {:?} found", hash); + let _ = find_and_send_tags( + opt.finish_tag.as_ref(), + &handler, + &pixiv, + &service, + &tmpdir, + &mut file, + ) + .await; + } + let elapsed = start.elapsed(); + + if elapsed.as_secs() < 8 { + tokio::time::sleep(sleep_duration - elapsed).await; // rate limit of 6# / 30s } - file.associate_urls(vec![url.to_string()]).await?; - } else { - tracing::info!("No pixiv post for file {:?} found", hash); } Ok(()) diff --git a/src/operations/find_and_send_tags.rs b/src/operations/find_and_send_tags.rs new file mode 100644 index 0000000..7d621de --- /dev/null +++ b/src/operations/find_and_send_tags.rs @@ -0,0 +1,65 @@ +use crate::{ + error::Result, + utils::pixiv::{get_pixiv_url, get_sauces_for_file, get_tags_for_sauce}, +}; +use hydrus_api::wrapper::{hydrus_file::HydrusFile, service::ServiceName}; +use pixiv_rs::PixivClient; +use rustnao::{Handler, Sauce}; +use tempdir::TempDir; + +pub async fn find_and_send_tags( + finish_tag: Option<&String>, + handler: &Handler, + pixiv: &PixivClient, + service: &ServiceName, + tmpdir: &TempDir, + mut file: &mut HydrusFile, +) -> Result<()> { + if let Err(e) = search_and_assign_tags(&handler, &pixiv, &service, &tmpdir, &mut file).await { + let hash = file.hash().await.unwrap(); + tracing::error!("Failed to search tags to file {}: {:?}", hash, e); + } else if let Some(finish_tag) = finish_tag { + file.add_tags(service.clone().into(), vec![finish_tag.into()]) + .await + .unwrap(); + } + + Ok(()) +} + +async fn search_and_assign_tags( + handler: &Handler, + pixiv: &PixivClient, + service: &ServiceName, + tmpdir: &TempDir, + mut file: &mut HydrusFile, +) -> Result<()> { + tracing::debug!("Getting tags for hydrus file {:?}", file.id); + let sauces = get_sauces_for_file(&handler, tmpdir, file).await?; + + assign_pixiv_tags_and_url(&pixiv, service, &mut file, &sauces).await +} + +async fn assign_pixiv_tags_and_url( + pixiv: &&PixivClient, + service: &ServiceName, + file: &mut &mut HydrusFile, + sauce: &Vec, +) -> Result<()> { + let hash = file.hash().await?; + if let Some(url) = get_pixiv_url(&sauce) { + let tags = get_tags_for_sauce(&pixiv, url).await?; + + if tags.len() > 0 { + tracing::info!("Found {} tags for file {:?}", tags.len(), hash); + file.add_tags(service.clone().into(), tags).await?; + } else { + tracing::info!("No tags for file {:?} found", hash); + } + file.associate_urls(vec![url.to_string()]).await?; + } else { + tracing::info!("No pixiv post for file {:?} found", hash); + } + + Ok(()) +} diff --git a/src/operations/find_and_send_urls.rs b/src/operations/find_and_send_urls.rs new file mode 100644 index 0000000..9a36843 --- /dev/null +++ b/src/operations/find_and_send_urls.rs @@ -0,0 +1,21 @@ +use hydrus_api::{wrapper::hydrus_file::HydrusFile, Hydrus}; +use rustnao::Handler; +use tempdir::TempDir; + +use crate::error::Result; +use crate::utils::pixiv::{get_sauces_for_file, get_urls}; + +pub async fn find_and_send_urls( + hydrus: &Hydrus, + handler: &Handler, + tmpdir: &TempDir, + file: &mut HydrusFile, +) -> Result<()> { + let sauces = get_sauces_for_file(&handler, tmpdir, file).await?; + let urls = get_urls(&sauces); + for url in urls { + hydrus.import().url(url).run().await?; + } + + Ok(()) +} diff --git a/src/operations/mod.rs b/src/operations/mod.rs new file mode 100644 index 0000000..14fb073 --- /dev/null +++ b/src/operations/mod.rs @@ -0,0 +1,2 @@ +pub mod find_and_send_tags; +pub mod find_and_send_urls; diff --git a/src/utils/mod.rs b/src/utils/mod.rs new file mode 100644 index 0000000..95543b2 --- /dev/null +++ b/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod pixiv; diff --git a/src/search.rs b/src/utils/pixiv.rs similarity index 63% rename from src/search.rs rename to src/utils/pixiv.rs index 28361b4..a95a26a 100644 --- a/src/search.rs +++ b/src/utils/pixiv.rs @@ -1,11 +1,26 @@ +use crate::error::Result; use hydrus_api::wrapper::hydrus_file::HydrusFile; use hydrus_api::wrapper::tag::Tag; use pixiv_rs::PixivClient; -use rustnao::Sauce; +use rustnao::{Handler, Sauce}; use std::fs; use std::path::PathBuf; use tempdir::TempDir; +pub async fn get_sauces_for_file( + handler: &Handler, + tmpdir: &TempDir, + mut file: &mut HydrusFile, +) -> Result> { + tracing::debug!("Creating tmp file for hydrus file {:?}", file.id); + let path = create_tmp_sauce_file(&tmpdir, &mut file).await?; + tracing::debug!("Getting sauce for hydrus file {:?}", file.id); + + let sauce = handler.get_sauce(path.to_str().unwrap(), None, None)?; + tracing::debug!("Getting tags for hydrus file {:?}", file.id); + Ok(sauce) +} + pub fn get_urls(sauce: &Vec) -> Vec<&String> { sauce.iter().flat_map(|s| &s.ext_urls).collect() } @@ -26,10 +41,7 @@ pub async fn get_tags_for_sauce(pixiv: &PixivClient, url: &String) -> crate::Res Ok(tags) } -pub async fn create_tmp_sauce_file( - tmpdir: &TempDir, - file: &mut HydrusFile, -) -> crate::Result { +async fn create_tmp_sauce_file(tmpdir: &TempDir, file: &mut HydrusFile) -> crate::Result { let hash = file.hash().await?; let bytes = file.retrieve().await?.bytes; let path = tmpdir.path().join(&hash);