Support multiple boorus

reverse_image_search/app.py

@@ -14,7 +14,7 @@
 from reverse_image_search.engines.saucenao import SauceNaoSearchEngine
 from reverse_image_search.providers import initiate_data_providers
 from reverse_image_search.providers.base import MessageConstruct
-from reverse_image_search.providers.danbooru import DanbooruProvider
+from reverse_image_search.providers.booru import BooruProvider
 from reverse_image_search.utils import chunks, download_file
 
 ZWS = "​"
@@ -25,7 +25,7 @@ class Arguments(Application.Arguments):
         downloads: Path
         file_url: str
         saucenao: SauceNaoSearchEngine.Config
-        danbooru: DanbooruProvider.Config
+        boorus: BooruProvider.Config
 
     arguments: "ReverseImageSearch.Arguments"
 
@@ -78,6 +78,7 @@ async def hndl_search(self, update: Update, _: ContextTypes.DEFAULT_TYPE):
             return
 
         file_url = self.arguments.file_url + file.name
+        file_url = "https://v2.nachtalb.io/ris/f/AQADlrsxG08-IVJ-.jpg"
 
         buttons = [
             InlineKeyboardButton(engine.name, engine.generate_search_url(str(file_url))) for engine in self.engines

reverse_image_search/engines/saucenao.py

@@ -110,5 +110,7 @@ async def search(self, file_url: str) -> AsyncGenerator[MessageConstruct, None]:
                 yield msg
 
     async def _booru(self, data: dict[str, dict[str, str | int | list[str]]]) -> MessageConstruct | None:
-        if danbooru_id := data["data"].get("danbooru_id"):
-            return await self._safe_search({"danbooru_id": danbooru_id}, "danbooru")
+        if post_id := data["data"].get("danbooru_id"):
+            return await self._safe_search({"id": post_id, "provider": "danbooru"}, "booru")
+        elif post_id := data["data"].get("yandere_id"):
+            return await self._safe_search({"id": post_id, "provider": "danbooru"}, "booru")

reverse_image_search/providers/__init__.py

@@ -1,8 +1,9 @@
 from typing import TYPE_CHECKING
+
 from aiohttp import ClientSession
 
 from reverse_image_search.providers.base import Provider
-from reverse_image_search.providers.danbooru import DanbooruProvider
+from reverse_image_search.providers.booru import BooruProvider
 
 if TYPE_CHECKING:
     from reverse_image_search.app import ReverseImageSearch
@@ -12,5 +13,5 @@ async def initiate_data_providers(
     session: ClientSession, config: "ReverseImageSearch.Arguments"
 ) -> dict[str, Provider]:
     return {
-        "danbooru": DanbooruProvider(session, config.danbooru.api_key, config.danbooru.username),
+        "booru": BooruProvider(session, config.boorus),
     }

reverse_image_search/providers/booru.py

@@ -0,0 +1,102 @@
+from typing import Any
+
+from aiohttp import BasicAuth, ClientSession
+from emoji import emojize
+from pydantic import BaseModel
+from tgtools.api.danbooru import DanbooruApi
+from tgtools.api.yandere import YandereApi
+from tgtools.models.booru_post import RATING
+from tgtools.models.danbooru_post import DanbooruPost
+from tgtools.telegram.text import tagified_string
+
+from .base import Info, MessageConstruct, Provider
+
+
+class BooruProvider(Provider):
+    """A provider for fetching and processing booru posts.
+
+    Attributes:
+        danbooru (DanbooruApi): An instance of the DanbooruApi class.
+        yandere (YandereApi): An instance of the YandereApi class.
+    """
+
+    class Config:
+        """Configuration for the BooruProvider.
+
+        Attributes:
+            danbooru_username (str): The username for accessing the Danbooru API.
+            danbooru_api_key (str): The API key for accessing the Danbooru API.
+        """
+
+        danbooru_username: str
+        danbooru_api_key: str
+
+    def __init__(self, session: ClientSession, config: "Config") -> None:
+        """
+        Initialise the BooruProvider with a session and configuration.
+
+        Args:
+            session (ClientSession): The aiohttp ClientSession to be used for API calls.
+            config (Config): The configuration object containing API credentials.
+        """
+        self.danbooru = DanbooruApi(session, BasicAuth(config.danbooru_username, config.danbooru_api_key))
+        self.yandere = YandereApi(session)
+
+    async def provide(self, data: dict[str, Any]) -> MessageConstruct | None:
+        """
+        Fetch and process a booru post.
+
+        Args:
+            data (dict[str, Any]): A dictionary containing the provider name and post ID.
+
+        Returns:
+            MessageConstruct | None: A MessageConstruct object containing the processed image
+                                     data or None if the provider is not supported.
+
+        Examples:
+            # Fetch a post from Danbooru with ID 12345
+            data = {"provider": "danbooru", "id": 12345}
+            message_construct = await booru_provider.provide(data)
+
+            # Fetch a post from Yandere with ID 67890
+            data = {"provider": "yandere", "id": 67890}
+            message_construct = await booru_provider.provide(data)
+        """
+        match data["provider"]:
+            case "danbooru":
+                provider = self.danbooru
+            case "yandere":
+                provider = self.yandere
+            case _:
+                return
+
+        post_id: int = data["id"]
+
+        post = await provider.post(post_id)
+
+        if post is None:
+            return
+
+        text: dict[str, str | Info] = {
+            "Rating": "",
+            "Artist": tagified_string(post.tags_artist),
+            "Tags": tagified_string(post.tags, 10),
+            "Characters": tagified_string(post.tags_character),
+            "Copyright": tagified_string(post.tags_copyright),
+        }
+        if post.rating:
+            rating_emoji = emojize(":no_one_under_eighteen:" if RATING.level(post.rating) > 1 else ":cherry_blossom:")
+            text["Rating"] = Info(f"{rating_emoji} {post.rating_simple}", "code")
+
+        result = MessageConstruct(
+            source_url=str(post.url),
+            additional_urls=[] if not post.source else [str(post.url)],
+            file=post.file_summary,
+            text=text,
+        )
+
+        match post:
+            case DanbooruPost():
+                result.source_url = f"https://www.pixiv.net/artworks/{post.pixiv_id}" if post.pixiv_id else post.source
+
+        return result