Add custom text for search result buttons

Nachtalb · Jul 24, 2023 · 83732e4 · 83732e4
1 parent f2159da
commit 83732e4
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 17 deletions.
diff --git a/reverse_image_search/app.py b/reverse_image_search/app.py
@@ -25,7 +25,6 @@
 from tgtools.telegram.compatibility import make_tg_compatible
 from tgtools.telegram.compatibility.base import OutputFileType
 from tgtools.utils.types import TELEGRAM_FILES
-from tgtools.utils.urls.emoji import FALLBACK_EMOJIS, host_name
 
 from reverse_image_search.engines import initiate_engines
 from reverse_image_search.engines.saucenao import SauceNaoSearchEngine
@@ -126,15 +125,7 @@ async def hndl_search(self, update: Update, _: ContextTypes.DEFAULT_TYPE) -> Non
     async def send_message_construct(
         self, result: SearchResult, query_message: Message, force_download: bool = False
     ) -> None:
-        buttons = [
-            InlineKeyboardButton(
-                host_name(result.message.provider_url, with_emoji=True, fallback=FALLBACK_EMOJIS["globe"]),
-                result.message.provider_url,
-            )
-        ]
-        for url in result.message.additional_urls:
-            buttons.append(InlineKeyboardButton(host_name(url, with_emoji=True), url=url))
-
+        buttons = [result.message.provider_button, *result.message.additional_buttons]
         markup = InlineKeyboardMarkup(tuple(chunks(buttons, 3)))
 
         additional_files_tasks = [

diff --git a/reverse_image_search/providers/base.py b/reverse_image_search/providers/base.py
@@ -3,7 +3,9 @@
 from typing import TYPE_CHECKING, Generic, Literal, Optional, Sequence, TypedDict, TypeVar, Union
 
 from pydantic import Field
+from telegram import InlineKeyboardButton
 from tgtools.models.summaries import Downloadable, FileSummary
+from tgtools.utils.urls.emoji import FALLBACK_EMOJIS, host_emoji, host_name
 
 if TYPE_CHECKING:
     from reverse_image_search.engines.base import SearchEngine
@@ -44,15 +46,15 @@ class MessageConstruct:
     """A data class representing a message construct.
 
     Attributes:
-        provider_url (str): The URL of the found media on the provider.
-        additional_urls (list[str]): A list of additional URLs related to the message.
+        provider_url (str | tuple[str, str]): The URL of the found media on the provider.
+        additional_urls (list[str | tuple[str, str]]): A list of additional URLs related to the message.
         text (dict[str, str | Info | None]): A dictionary containing the text elements of the message.
         file (MediaSummary | None): The main media file sent with the result message.
         additional_files (Sequence[MediaSummary] | None): Max 10 additional media files, sent in a media group
     """
 
-    provider_url: str
-    additional_urls: list[str]
+    provider_url: str | tuple[str, str]
+    additional_urls: list[str | tuple[str, str]]
     text: dict[str, str | Info | None]
     file: Optional[Union[FileSummary, Downloadable]] = None
     additional_files: Sequence[Union[FileSummary, Downloadable]] = Field(default_factory=list)
@@ -62,6 +64,32 @@ class MessageConstruct:
     def caption(self) -> str:
         return "\n".join(f"<b>{title}:</b> {str(content)}" for title, content in self.text.items() if content)
 
+    @property
+    def provider_button(self) -> InlineKeyboardButton:
+        if isinstance(self.provider_url, str):
+            text = host_name(self.provider_url, with_emoji=True, fallback=FALLBACK_EMOJIS["globe"])
+            url = self.provider_url
+        else:
+            text, url = self.provider_url
+            emoji = host_emoji(url, fallback=FALLBACK_EMOJIS["globe"])
+            text = f"{emoji} {text}"
+        return InlineKeyboardButton(text=text, url=url)
+
+    @property
+    def additional_buttons(self) -> list[InlineKeyboardButton]:
+        buttons = []
+        for url_or_tuple in self.additional_urls:
+            if isinstance(url_or_tuple, str):
+                text = host_name(url_or_tuple, with_emoji=True, fallback=FALLBACK_EMOJIS["globe"])
+                url = url_or_tuple
+            else:
+                text, url = url_or_tuple
+                emoji = host_emoji(url, fallback=FALLBACK_EMOJIS["globe"])
+                text = f"{emoji} {text}"
+            buttons.append(InlineKeyboardButton(text=text, url=url))
+
+        return buttons
+
 
 @dataclass
 class ProviderInfo:

diff --git a/reverse_image_search/providers/pixiv.py b/reverse_image_search/providers/pixiv.py
@@ -82,7 +82,6 @@ async def provide(self, data: PixivQuery) -> MessageConstruct | None:
             ),
         }
         source_url = f"https://www.pixiv.net/en/artworks/{post.id}"
-        artist_url = f"https://www.pixiv.net/en/user/{post.user.id}"
 
         client = post.get_client()
         main_file = ToDownload(
@@ -109,9 +108,9 @@ async def provide(self, data: PixivQuery) -> MessageConstruct | None:
                 ]
 
         return MessageConstruct(
-            provider_url=str(source_url),
+            provider_url=("Artwork", str(source_url)),
             additional_urls=[
-                artist_url,
+                ("Artist", str(source_url)),
             ],
             text=text,
             file=main_file,