Merge pull request #46 from nitori-dev/pixivtags

Pixiv Tagging

config_sample.toml

@@ -16,6 +16,7 @@ deepbooru_set_tag = true  # Tag Deepbooru post with tag "deepbooru"
 hide_progress = false  # Set this to true to hide the progress bar
 use_pixiv_artist = false  # If the artist could only be found on pixiv, create and use the pixiv artist. Category has to be 'artist'.
 update_relations = true  # Set character <> parody relation if SauceNAO is disabled (or limit reached) and Deepbooru enabled
+pixiv_tags = false  # Tag posts with pixiv tags
 
 # Leave those options as "None" if you have no accounts on these sites
 [danbooru]

pyproject.toml

@@ -83,6 +83,7 @@ gallery-dl = "^1.25.4"
 cunnypy = "^2.0.0"
 tldextract = "^3.4.4"
 flask = "^3.0.0"
+pixivpy3 = "^3.7.4"
 
 [tool.poetry.dev-dependencies]
 black = "^22.1.0"

src/szurubooru_toolkit/__init__.py

@@ -6,6 +6,7 @@
 from .config import Config
 from .danbooru import Danbooru  # noqa F401
 from .gelbooru import Gelbooru  # noqa F401
+from .pixiv import Pixiv
 from .szurubooru import Post  # noqa F401
 from .szurubooru import Szurubooru
 from .twitter import Twitter  # noqa F401

src/szurubooru_toolkit/config.py

@@ -61,6 +61,7 @@ def check_attr_set(self) -> None:
                 'hide_progress',
                 'use_pixiv_artist',
                 'update_relations',
+                'pixiv_tags',
             ],
             'upload_media': [
                 'src_path',

src/szurubooru_toolkit/pixiv.py

@@ -0,0 +1,48 @@
+from pixivpy3 import AppPixivAPI as Pixiv_Module
+from aiohttp.client_exceptions import ClientConnectorError
+from loguru import logger
+
+class Pixiv:
+    def __init__(self, token):
+        self.client = Pixiv_Module()
+        self.client.auth(refresh_token=token)
+
+
+    def get_result(self, result_url):
+        temp = result_url.split('=')
+        post_id = int(result_url.split('=')[-1])
+        logger.debug(f'Getting result from id {post_id}')
+        for _ in range(1, 12):
+            try:
+                result = self.client.illust_detail(post_id)
+                logger.debug(f'Returning result: {result}')
+                break
+            except ClientConnectorError:
+                logger.debug('Could not establish connection to Pixiv, trying again in 5s...')
+                sleep(5)
+            except KeyError:  # In case the post got deleted but is still indexed
+                result = None
+                logger.debug('Got no result')
+                break
+        else:
+            result = None
+
+        return result
+
+    def get_tags(self, result):
+        tags = []
+        if result.illust and result.illust.tags:
+            for tag in result.illust.tags:
+                temp = tag['name']
+                if not temp == None:
+                    if not temp == "R-18":
+                        tags.append(temp)
+        logger.debug(f'Returning tags {tags}')
+        return tags
+
+    def get_rating(self, result):
+        if result.illust and result.illust.tags:
+            for tag in result.illust.tags:
+                if tag['name'] == "R-18":
+                    return 'unsafe'
+        return 'safe'

src/szurubooru_toolkit/scripts/auto_tagger.py

@@ -227,7 +227,7 @@ def main(post_id: str = None, file_to_upload: bytes = None, limit_reached: bool
                     md5_results = asyncio.run(search_boorus('all', 'md5:' + post.md5, 1, 0))
 
                 if md5_results:
-                    tags_by_md5, sources, post.rating = prepare_post(md5_results)
+                    tags_by_md5, sources, post.rating = prepare_post(md5_results, config)
                     post.source = collect_sources(*sources, *post.source.splitlines())
                 else:
                     tags_by_md5 = []
@@ -257,7 +257,7 @@ def main(post_id: str = None, file_to_upload: bytes = None, limit_reached: bool
                 sauce_results, limit_reached = get_saucenao_results(sauce, post, image)
 
                 if sauce_results:
-                    tags_by_sauce, sources, post.rating = prepare_post(sauce_results)
+                    tags_by_sauce, sources, post.rating = prepare_post(sauce_results, config)
                     post.source = collect_sources(*sources, *post.source.splitlines())
                 else:
                     tags_by_sauce = []

src/szurubooru_toolkit/utils.py

@@ -26,6 +26,7 @@
 from szurubooru_toolkit import Config
 from szurubooru_toolkit import Danbooru
 from szurubooru_toolkit import Gelbooru
+from szurubooru_toolkit import Pixiv
 
 
 # Keep track of total tagged posts
@@ -513,25 +514,36 @@ async def search_boorus(booru: str, query: str, limit: int, page: int = 1) -> di
     return results
 
 
-def prepare_post(results: dict):
+def prepare_post(results: dict, config: Config):
     tags = []
     sources = []
     rating = []
-
+    booru_found = False
     for booru, result in results.items():
         if booru != 'pixiv':
             tags.append(result[0].tags.split())
             sources.append(generate_src({'site': booru, 'id': result[0].id}))
             rating = convert_rating(result[0].rating)
+            booru_found = True
         else:
+            if config.auto_tagger['pixiv_tags']:
+                pixiv = Pixiv(config.pixiv['token'])
+                pixiv_result = pixiv.get_result(results['pixiv'].url)
+                pixiv_tags = pixiv.get_tags(pixiv_result)
+                pixiv_rating = pixiv.get_rating(pixiv_result)
             pixiv_sources, pixiv_artist = extract_pixiv_artist(results['pixiv'])
             sources.append(pixiv_sources)
 
     final_tags = [item for sublist in tags for item in sublist]
-
-    if not final_tags and 'pixiv' in results and pixiv_artist:
-        final_tags = pixiv_artist
-
+    if config.auto_tagger['pixiv_tags']:
+        if 'pixiv' in results and pixiv_tags:
+            final_tags.extend(pixiv_tags)
+        if not booru_found and 'pixiv' in results and pixiv_rating:
+            rating = pixiv_rating
+            final_tags.append('check_safety')
+
+    if not booru_found and 'pixiv' in results and pixiv_artist:
+        final_tags.append(pixiv_artist)
     return final_tags, sources, rating
 
 
@@ -565,7 +577,7 @@ def extract_pixiv_artist(result: Any) -> tuple[str, list]:
     else:
         artist = None
 
-    return source, [artist]
+    return source, artist
 
 
 def extract_twitter_artist(metadata: dict) -> str: