Fix rate limit and add additional tag for marking of files

Fix handling of the rate limit. Also errors are now propagated instead
of panicing when one occurs. With the additional tag --finish-tag
one can assign a tag to all files that have been processed.

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -598,7 +598,7 @@ dependencies = [
 
 [[package]]
 name = "hydrus-pixiv-tagger"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
  "env_logger",
  "hydrus-api",
@@ -607,6 +607,7 @@ dependencies = [
  "rustnao",
  "structopt",
  "tempdir",
+ "thiserror",
  "tokio 1.9.0",
 ]
 
@@ -1669,6 +1670,26 @@ dependencies = [
  "unicode-width",
 ]
 
+[[package]]
+name = "thiserror"
+version = "1.0.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "93119e4feac1cbe6c798c34d3a53ea0026b0b1de6a120deef895137c0529bfe2"
+dependencies = [
+ "thiserror-impl",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "060d69a0afe7796bf42e9e2ff91f5ee691fb15c53d38b4b62a9a53eb23164745"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "time"
 version = "0.1.43"

Cargo.toml

@@ -1,10 +1,10 @@
 [package]
 name = "hydrus-pixiv-tagger"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2018"
 license = "Apache-2.0"
 description = "Automatically tag hydrus file by using pixiv and saucenao"
-author = ["trivernis <trivernis@protonmail.com>"]
+authors = ["trivernis <trivernis@protonmail.com>"]
 readme = "README.md"
 repository = "https://github.com/Trivernis/hydrus-pixiv-tagger"
 
@@ -18,6 +18,7 @@ structopt = "0.3.22"
 tempdir = "0.3.7"
 env_logger = "0.9.0"
 log = "0.4.14"
+thiserror = "1.0.26"
 
 [dependencies.tokio]
 version = "1.9.0"

README.md

@@ -30,12 +30,12 @@ FLAGS:
     -V, --version    Prints version information
 
 OPTIONS:
+        --finish-tag <finish-tag>        Tag that is assigned to files that have been processed
         --hydrus-key <hydrus-key>        The hydrus client api key
         --hydrus-url <hydrus-url>        The url to the hydrus client api [default: http://127.0.0.1:45869]
         --saucenao-key <saucenao-key>    The saucenao api key
         --tag-service <tag-service>      The tag service the tags will be assigned to [default: my tags]
     -t, --tags <tags>...                 Tags used to search for files
-
 ```
 
 ## Example
@@ -56,6 +56,7 @@ hydrus-pixiv-tagger
     --saucenao-key <key2>\
     --inbox
     --tag-service 'my tags'
+    --finish-tag 'meta:automatically_tagged'
 ```
 
 ## License

src/error.rs

@@ -0,0 +1,27 @@
+use hydrus_api::error::Error as HydrusError;
+use pixiv_rs::error::Error as PixivError;
+use rustnao::Error as RustNaoError;
+use thiserror::Error;
+
+pub type Result<T> = std::result::Result<T, Error>;
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error(transparent)]
+    Pixiv(#[from] PixivError),
+
+    #[error("{0}")]
+    RustNao(String),
+
+    #[error(transparent)]
+    Hydrus(#[from] HydrusError),
+
+    #[error(transparent)]
+    Io(#[from] std::io::Error),
+}
+
+impl From<RustNaoError> for Error {
+    fn from(e: RustNaoError) -> Self {
+        Self::RustNao(e.to_string())
+    }
+}

src/main.rs

@@ -1,3 +1,6 @@
+mod error;
+
+use crate::error::Result;
 use hydrus_api::api_core::searching_and_fetching_files::FileSearchLocation;
 use hydrus_api::wrapper::hydrus_file::HydrusFile;
 use hydrus_api::wrapper::service::ServiceName;
@@ -37,6 +40,10 @@ struct Opt {
     /// Searches in the inbox instead
     #[structopt(long)]
     inbox: bool,
+
+    /// Tag that is assigned to files that have been processed
+    #[structopt(long)]
+    finish_tag: Option<String>,
 }
 
 #[tokio::main]
@@ -64,11 +71,19 @@ async fn main() {
     let files = hydrus.search(search_location, tags).await.unwrap();
     let tmpdir = TempDir::new("hydrus-files").unwrap();
 
-    let sleep_duration = Duration::from_secs(5);
+    let sleep_duration = Duration::from_secs(6);
 
     for mut file in files {
         let start = Instant::now();
-        search_and_assign_tags(&handler, &pixiv, &service, &tmpdir, &mut file).await;
+        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();
+        }
         let elapsed = start.elapsed();
 
         if elapsed.as_secs() < 6 {
@@ -83,14 +98,12 @@ async fn search_and_assign_tags(
     service: &ServiceName,
     tmpdir: &TempDir,
     mut file: &mut HydrusFile,
-) {
+) -> Result<()> {
     log::debug!("Creating tmp file for hydrus file {:?}", file.id);
-    let path = create_tmp_sauce_file(&tmpdir, &mut file).await;
+    let path = create_tmp_sauce_file(&tmpdir, &mut file).await?;
     log::debug!("Getting sauce for hydrus file {:?}", file.id);
 
-    let sauce = handler
+    let sauce = handler.get_sauce(path.to_str().unwrap(), None, None)?;
-        .get_sauce(path.to_str().unwrap(), None, None)
-        .unwrap();
     log::debug!("Getting tags for hydrus file {:?}", file.id);
 
     assign_pixiv_tags_and_url(&pixiv, service, &mut file, &sauce).await
@@ -105,24 +118,31 @@ async fn assign_pixiv_tags_and_url(
     service: &ServiceName,
     file: &mut &mut HydrusFile,
     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;
+        let tags = get_tags_for_sauce(&pixiv, url).await?;
 
         if tags.len() > 0 {
-            log::info!("Found {} tags for file {:?}", tags.len(), file.id);
+            log::info!("Found {} tags for file {:?}", tags.len(), hash);
-            file.add_tags(service.clone(), tags).await.unwrap();
+            file.add_tags(service.clone(), tags).await?;
+        } else {
+            log::info!("No tags for file {:?} found", hash);
         }
-        file.associate_urls(vec![url.to_string()]).await.unwrap();
+        file.associate_urls(vec![url.to_string()]).await?;
+    } else {
+        log::info!("No pixiv post for file {:?} found", hash);
     }
+
+    Ok(())
 }
 
-async fn get_tags_for_sauce(pixiv: &PixivClient, url: &String) -> Vec<Tag> {
+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.unwrap();
+        let illustration = pixiv.illustration(pixiv_id).await?;
 
         for tag in illustration.tags.tags {
             let tag_value = tag.translation.get("en").unwrap_or(&tag.tag);
@@ -130,13 +150,14 @@ async fn get_tags_for_sauce(pixiv: &PixivClient, url: &String) -> Vec<Tag> {
         }
     }
 
-    tags
+    Ok(tags)
 }
 
-async fn create_tmp_sauce_file(tmpdir: &TempDir, file: &mut HydrusFile) -> PathBuf {
+async fn create_tmp_sauce_file(tmpdir: &TempDir, file: &mut HydrusFile) -> Result<PathBuf> {
-    let hash = file.hash().await.unwrap();
+    let hash = file.hash().await?;
-    let bytes = file.retrieve().await.unwrap().bytes;
+    let bytes = file.retrieve().await?.bytes;
     let path = tmpdir.path().join(&hash);
-    fs::write(&path, bytes).unwrap();
+    fs::write(&path, bytes)?;
-    path
+
+    Ok(path)
 }