Add command to send all found urls instead of pixiv tags

Signed-off-by: trivernis <trivernis@protonmail.com>
main
trivernis 2 years ago
parent 99109eaca1
commit 4fdba3c6a8
Signed by: Trivernis
GPG Key ID: DFFFCC2C7A02DB45

602
Cargo.lock generated

File diff suppressed because it is too large Load Diff

@ -12,14 +12,14 @@ repository = "https://github.com/Trivernis/hydrus-pixiv-tagger"
[dependencies]
pixiv-rs = "0.1.0"
hydrus-api = "0.6.0"
rustnao = "0.2"
structopt = "0.3.22"
hydrus-api = "0.9.1"
rustnao = "0.2.1"
structopt = "0.3.26"
tempdir = "0.3.7"
env_logger = "0.9.0"
log = "0.4.14"
thiserror = "1.0.26"
log = "0.4.16"
thiserror = "1.0.30"
[dependencies.tokio]
version = "1.9.0"
features = ["macros", "rt", "time"]
version = "1.17.0"
features = ["macros", "rt", "time"]

@ -1,22 +1,31 @@
mod error;
pub mod search;
use crate::error::Result;
use crate::search::get_urls;
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 std::fs;
use std::path::PathBuf;
use structopt::clap::AppSettings;
use structopt::StructOpt;
use tempdir::TempDir;
use tokio::time::{Duration, Instant};
#[derive(StructOpt, Debug)]
#[structopt()]
enum Command {
/// Sends urls to hydrus to be imported
SendUrl(Options),
/// Maps the tags found for the hydrus url to the hydrus file
SendTags(Options),
}
#[derive(StructOpt, Debug, Clone)]
#[structopt(settings = &[AppSettings::AllowLeadingHyphen])]
struct Opt {
struct Options {
/// The saucenao api key
#[structopt(long, env)]
saucenao_key: String,
@ -45,7 +54,11 @@ struct Opt {
#[tokio::main(flavor = "current_thread")]
async fn main() {
env_logger::builder().init();
let opt: Opt = Opt::from_args();
let cmd: Command = Command::from_args();
let opt = match &cmd {
Command::SendUrl(opt) => opt.clone(),
Command::SendTags(opt) => opt.clone(),
};
let handler = HandlerBuilder::new()
.api_key(&opt.saucenao_key)
@ -67,14 +80,21 @@ async fn main() {
for mut file in files {
let start = Instant::now();
if let Err(e) = search_and_assign_tags(&handler, &pixiv, &service, &tmpdir, &mut file).await
{
let hash = file.hash().await.unwrap();
log::error!("Failed to search and assign tags to file {}: {:?}", hash, e);
} else if let Some(finish_tag) = &opt.finish_tag {
file.add_tags(service.clone(), vec![finish_tag.into()])
.await
.unwrap();
match &cmd {
Command::SendUrl(_) => {
let _ = search_and_send_urls(&hydrus, &handler, &tmpdir, &mut file).await;
}
Command::SendTags(_) => {
let _ = tag_file(
opt.finish_tag.as_ref(),
&handler,
&pixiv,
&service,
&tmpdir,
&mut file,
)
.await;
}
}
let elapsed = start.elapsed();
@ -84,6 +104,41 @@ async fn main() {
}
}
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();
log::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,
@ -91,18 +146,24 @@ async fn search_and_assign_tags(
tmpdir: &TempDir,
mut file: &mut HydrusFile,
) -> Result<()> {
log::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 get_sauces_for_file(
handler: &Handler,
tmpdir: &TempDir,
mut file: &mut HydrusFile,
) -> Result<Vec<Sauce>> {
log::debug!("Creating tmp file for hydrus file {:?}", file.id);
let path = create_tmp_sauce_file(&tmpdir, &mut file).await?;
let path = search::create_tmp_sauce_file(&tmpdir, &mut file).await?;
log::debug!("Getting sauce for hydrus file {:?}", file.id);
let sauce = handler.get_sauce(path.to_str().unwrap(), None, None)?;
log::debug!("Getting tags for hydrus file {:?}", file.id);
assign_pixiv_tags_and_url(&pixiv, service, &mut file, &sauce).await
}
fn get_pixiv_url(sauce: &Vec<Sauce>) -> Option<&String> {
sauce.first().and_then(|s| s.ext_urls.first())
Ok(sauce)
}
async fn assign_pixiv_tags_and_url(
@ -112,12 +173,12 @@ async fn assign_pixiv_tags_and_url(
sauce: &Vec<Sauce>,
) -> Result<()> {
let hash = file.hash().await?;
if let Some(url) = get_pixiv_url(&sauce) {
let tags = get_tags_for_sauce(&pixiv, url).await?;
if let Some(url) = search::get_pixiv_url(&sauce) {
let tags = search::get_tags_for_sauce(&pixiv, url).await?;
if tags.len() > 0 {
log::info!("Found {} tags for file {:?}", tags.len(), hash);
file.add_tags(service.clone(), tags).await?;
file.add_tags(service.clone().into(), tags).await?;
} else {
log::info!("No tags for file {:?} found", hash);
}
@ -128,28 +189,3 @@ async fn assign_pixiv_tags_and_url(
Ok(())
}
async fn get_tags_for_sauce(pixiv: &PixivClient, url: &String) -> Result<Vec<Tag>> {
let mut tags = Vec::new();
if let Some(pixiv_id) = url.rsplit_once("=").map(|s| s.1) {
log::trace!("Pixiv id is '{}'", pixiv_id);
let illustration = pixiv.illustration(pixiv_id).await?;
for tag in illustration.tags.tags {
let tag_value = tag.translation.get("en").unwrap_or(&tag.tag);
tags.push(Tag::from(tag_value));
}
}
Ok(tags)
}
async fn create_tmp_sauce_file(tmpdir: &TempDir, file: &mut HydrusFile) -> Result<PathBuf> {
let hash = file.hash().await?;
let bytes = file.retrieve().await?.bytes;
let path = tmpdir.path().join(&hash);
fs::write(&path, bytes)?;
Ok(path)
}

@ -0,0 +1,43 @@
use hydrus_api::wrapper::hydrus_file::HydrusFile;
use hydrus_api::wrapper::tag::Tag;
use pixiv_rs::PixivClient;
use rustnao::Sauce;
use std::fs;
use std::path::PathBuf;
use tempdir::TempDir;
pub fn get_urls(sauce: &Vec<Sauce>) -> Vec<&String> {
sauce.iter().flat_map(|s| &s.ext_urls).collect()
}
pub async fn get_tags_for_sauce(pixiv: &PixivClient, url: &String) -> crate::Result<Vec<Tag>> {
let mut tags = Vec::new();
if let Some(pixiv_id) = url.rsplit_once("=").map(|s| s.1) {
log::trace!("Pixiv id is '{}'", pixiv_id);
let illustration = pixiv.illustration(pixiv_id).await?;
for tag in illustration.tags.tags {
let tag_value = tag.translation.get("en").unwrap_or(&tag.tag);
tags.push(Tag::from(tag_value));
}
}
Ok(tags)
}
pub async fn create_tmp_sauce_file(
tmpdir: &TempDir,
file: &mut HydrusFile,
) -> crate::Result<PathBuf> {
let hash = file.hash().await?;
let bytes = file.retrieve().await?.bytes;
let path = tmpdir.path().join(&hash);
fs::write(&path, bytes)?;
Ok(path)
}
pub fn get_pixiv_url(sauce: &Vec<Sauce>) -> Option<&String> {
sauce.first().and_then(|s| s.ext_urls.first())
}
Loading…
Cancel
Save