From 1c779e98a89bbb612a8dd86a83e98a2775ae4cf1 Mon Sep 17 00:00:00 2001 From: YoursFunny Date: Mon, 24 Jun 2024 01:25:03 +0800 Subject: [PATCH] add pixiv support --- main.py | 10 ++- utils/pixiv.py | 199 ++++++++++++++++++++++++++++++++++++++++++++++ utils/regex.py | 4 +- utils/telegram.py | 6 +- 4 files changed, 214 insertions(+), 5 deletions(-) create mode 100644 utils/pixiv.py diff --git a/main.py b/main.py index 1fc5e2c..9a500be 100644 --- a/main.py +++ b/main.py @@ -39,8 +39,7 @@ async def inline_query(update: Update, context: ContextTypes.DEFAULT_TYPE) -> No return logger.info(f"Query: {query}") async with Telegram(query) as tweet: - result = list(tweet.inline_query_generator) - await update.inline_query.answer(result) + await update.inline_query.answer(tweet.inline_query_result()) @send_action(ChatAction.UPLOAD_PHOTO) @@ -48,7 +47,10 @@ async def url_media(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: url = update.message.text logger.info(f"Receiving url: {url}") async with Telegram(url) as tweet: - media = list(tweet.message_media_generator) + media = tweet.message_media_result() + if not media: + await update.effective_message.reply_text("No media found or media type is not supported.") + return message_to_send = await update.effective_message.reply_media_group( media, caption=tweet.message_text, @@ -207,6 +209,8 @@ async def post_init(application: Application) -> None: await application.bot.set_my_description(DESCRIPTION) await application.bot.set_my_short_description(DESCRIPTION) Telegram.init_client() + if common.PIXIV_REFRESH_TOKEN: + await Pixiv.init_client(common.PIXIV_REFRESH_TOKEN) async def post_stop(application: Application) -> None: diff --git a/utils/pixiv.py b/utils/pixiv.py new file mode 100644 index 0000000..cbe545e --- /dev/null +++ b/utils/pixiv.py @@ -0,0 +1,199 @@ +import html +from typing import Generator, Literal, TYPE_CHECKING +from uuid import uuid4 + +from async_pixiv import PixivClient +from telegram import InlineQueryResultPhoto, InputMediaPhoto + +from utils.logger import get_logger + +if TYPE_CHECKING: + from async_pixiv.model.illust import Illust + from utils.types import TypeInlineQueryResult, TypeMessageMediaResult + +logger = get_logger(__name__) + +message_raw_text = """{text} / {author} +{tags} +""" + + +class PixivMedia: + __slots__ = ('_url', '_thumb') + + def __init__(self, url: str, thumb: str): + self._url: str = url + self._thumb: str = thumb + + def __str__(self): + return f"PixivMedia(url={self.url}, thumb={self.thumb}, large={self.large})" + + @property + def url(self) -> str: + return self._url + + @property + def thumb(self) -> str: + return self._thumb + + @property + def large(self) -> str: + url = self._url.replace("img-original", "img-master").removesuffix(".jpg").removesuffix(".png") + return url + "_master1200.jpg" + + +class Pixiv: + __slots__ = ('_illust',) + + def __init__(self, illust: Illust): + self._illust: Illust = illust + + @property + def url(self) -> str: + return str(self._illust.link).rstrip('/') + + @property + def type(self) -> Literal["illust", "manga", "ugoira"]: + return self._illust.type.value + + @property + def title(self) -> str: + return self._illust.title + + @property + def author(self) -> str: + return self._illust.user.name + + @property + def author_url(self) -> str: + return str(self._illust.user.link) + + @property + def description(self) -> str: + return self._illust.caption + + @property + def tags(self) -> list[str]: + return [tag.name for tag in self._illust.tags] + + @property + def is_multiple_pages(self) -> bool: + return self._illust.page_count > 1 + + @property + def is_nsfw(self) -> bool: + return self._illust.is_nsfw + + @property + def is_ai(self) -> bool: + return self._illust.ai_type.value == 2 + + @property + def images(self) -> list[PixivMedia]: + if self.is_multiple_pages: + return [ + PixivMedia( + url=page.image_urls.original, + thumb=page.image_urls.medium + ) + for page in self._illust.meta_pages + ] + else: + return [ + PixivMedia( # should observe if image_url.original is always None for single page + url=self._illust.image_urls.original or self._illust.meta_single_page.original, + thumb=self._illust.image_urls.medium + ) + ] + + +class ProcessPixiv: + _client: PixivClient + + __slots__ = ('_url', '_illust') + + @classmethod + async def init_client(cls, token: str) -> None: + cls._client = PixivClient() + await cls._client.login_with_token(token) + + @classmethod + async def close_client(cls) -> None: + await cls._client.close() + + def __init__(self, url: str): + self._url: str = url + + async def __aenter__(self): + self._illust = await self._fetch_illust() + return Pixiv(self._illust) + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + async def _fetch_illust(self) -> Illust: + illust_id = self._parse_illust_id() + return (await self._client.ILLUST.detail(illust_id)).illust + + def _parse_illust_id(self) -> int: + return int(self._url.split("/")[-1]) + + +class TelegramPixiv: + __slots__ = ('_url', '_pixiv') + + def __init__(self, url: str): + self._url = url + + async def __aenter__(self): + async with ProcessPixiv(self._url) as pixiv: + self._pixiv = pixiv + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + pass + + def url(self) -> str: + return self._url + + @property + def message_text(self) -> str: + pixiv = self._pixiv + return message_raw_text.format( + url=pixiv.url, + author_url=pixiv.author_url, + author=html.escape(pixiv.author), + text=html.escape(pixiv.title), + tags=html.escape(" ".join(f"#{name}" for name in pixiv.tags)) + ) + + def inline_query_result(self) -> tuple[TypeInlineQueryResult, ...]: + return tuple(self.inline_query_generator()) + + def message_media_result(self) -> tuple[TypeMessageMediaResult, ...]: + return tuple(self.message_media_generator()) + + def inline_query_generator(self) -> Generator[TypeInlineQueryResult, None, None]: + pixiv = self._pixiv + for media in pixiv.images: + logger.info(str(media)) + if pixiv.type in ("illust", "manga"): + yield InlineQueryResultPhoto( + id=str(uuid4()), + photo_url=media.large, + thumbnail_url=media.thumb, + caption=self.message_text + ) + else: + yield + + def message_media_generator(self) -> Generator[TypeMessageMediaResult, None, None]: + pixiv = self._pixiv + for media in pixiv.images: + if pixiv.type in ("illust", "manga"): + yield InputMediaPhoto( + media=media.url, + has_spoiler=pixiv.is_nsfw + ) + else: + yield diff --git a/utils/regex.py b/utils/regex.py index a5b3144..389b4a6 100644 --- a/utils/regex.py +++ b/utils/regex.py @@ -1,6 +1,8 @@ import re -x_url = re.compile(r"^(?:https?://)(?:www\.|mobile\.|)(?:x|twitter|fixvx|vxtwitter)\.com/(.+)/status/(\d+)") +x_url = re.compile(r"^(?:https?://)?(?:www\.|mobile\.)?(?:x|twitter|fixvx|vxtwitter)\.com/(.+)/status/(\d+)") x_media_url = re.compile(r"^(?:https?://)(pbs|video)\.twimg\.com/(.*)") x_tco_url = re.compile(r"(?:https?://)t\.co/.+$", re.M) message_url = re.compile(r"\[.+]", re.S) + +pixiv_url = re.compile(r"^(?:https?://)?(?:www\.)?pixiv\.net/(?:en/)?artworks/(\d+)") diff --git a/utils/telegram.py b/utils/telegram.py index e00ec54..c060503 100644 --- a/utils/telegram.py +++ b/utils/telegram.py @@ -1,7 +1,8 @@ from __future__ import annotations from .net import NetClient -from .regex import x_url +from .pixiv import TelegramPixiv +from .regex import pixiv_url, x_url from .tweet import TelegramTweet @@ -13,6 +14,9 @@ class Telegram(NetClient): if x_url.match(self._url): async with TelegramTweet(self._url) as tweet: return tweet + elif pixiv_url.match(self._url): + async with TelegramPixiv(self._url) as pixiv: + return pixiv async def __aexit__(self, exc_type, exc_val, exc_tb): pass