from __future__ import annotations import html from typing import Generator, Literal, TYPE_CHECKING from uuid import uuid4 from async_pixiv import PixivClient from async_pixiv.error import ApiError 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).rstrip('/') @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) cls._token = token @classmethod async def close_client(cls) -> None: await cls._client.close() @classmethod async def refresh_token(cls): await cls._client.login_with_token(cls._token) 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() try: return (await self._client.ILLUST.detail(illust_id)).illust except ApiError: await self.refresh_token() return (await self._client.ILLUST.detail(illust_id)).illust # TODO use retry here 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 @property def url(self) -> str: return self._pixiv.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: logger.info(str(media)) if pixiv.type in ("illust", "manga"): yield InputMediaPhoto( media=media.large, has_spoiler=pixiv.is_nsfw ) else: yield