PamGram/core/template/models.py

from enum import Enum
from typing import Optional, Union, List

from telegram import Message, InputMediaPhoto, InputMediaDocument

from core.template.cache import HtmlToFileIdCache
from core.template.error import ErrorFileType, FileIdNotFound


class FileType(Enum):
    PHOTO = 1
    DOCUMENT = 2

    @staticmethod
    def media_type(file_type: "FileType"):
        """对应的 Telegram media 类型"""
        if file_type == FileType.PHOTO:
            return InputMediaPhoto
        elif file_type == FileType.DOCUMENT:
            return InputMediaDocument
        else:
            raise ErrorFileType


class RenderResult:
    """渲染结果"""

    def __init__(
        self,
        html: str,
        photo: Union[bytes, str],
        file_type: FileType,
        cache: HtmlToFileIdCache,
        ttl: int = 24 * 60 * 60,
        caption: Optional[str] = None,
        parse_mode: Optional[str] = None,
        filename: Optional[str] = None,
    ):
        """
        `html`: str 渲染生成的 html
        `photo`: Union[bytes, str] 渲染生成的图片。bytes 表示是图片，str 则为 file_id
        """
        self.caption = caption
        self.parse_mode = parse_mode
        self.filename = filename
        self.html = html
        self.photo = photo
        self.file_type = file_type
        self._cache = cache
        self.ttl = ttl

    async def reply_photo(self, message: Message, *args, **kwargs):
        """是 `message.reply_photo` 的封装，上传成功后，缓存 telegram 返回的 file_id，方便重复使用"""
        if self.file_type != FileType.PHOTO:
            raise ErrorFileType

        reply = await message.reply_photo(photo=self.photo, *args, **kwargs)

        await self.cache_file_id(reply)

        return reply

    async def reply_document(self, message: Message, *args, **kwargs):
        """是 `message.reply_document` 的封装，上传成功后，缓存 telegram 返回的 file_id，方便重复使用"""
        if self.file_type != FileType.DOCUMENT:
            raise ErrorFileType

        reply = await message.reply_document(document=self.photo, *args, **kwargs)

        await self.cache_file_id(reply)

        return reply

    async def edit_media(self, message: Message, *args, **kwargs):
        """是 `message.edit_media` 的封装，上传成功后，缓存 telegram 返回的 file_id，方便重复使用"""
        if self.file_type != FileType.PHOTO:
            raise ErrorFileType

        media = InputMediaPhoto(
            media=self.photo, caption=self.caption, parse_mode=self.parse_mode, filename=self.filename
        )

        edit_media = await message.edit_media(media, *args, **kwargs)

        await self.cache_file_id(edit_media)

        return edit_media

    async def cache_file_id(self, reply: Message):
        """缓存 telegram 返回的 file_id"""
        if self.is_file_id():
            return

        if self.file_type == FileType.PHOTO and reply.photo:
            file_id = reply.photo[0].file_id
        elif self.file_type == FileType.DOCUMENT and reply.document:
            file_id = reply.document.file_id
        else:
            raise FileIdNotFound
        await self._cache.set_data(self.html, self.file_type.name, file_id, self.ttl)

    def is_file_id(self) -> bool:
        return isinstance(self.photo, str)


class RenderGroupResult:
    def __init__(self, results: List[RenderResult]):
        self.results = results

    async def reply_media_group(self, message: Message, *args, **kwargs):
        """是 `message.reply_media_group` 的封装，上传成功后，缓存 telegram 返回的 file_id，方便重复使用"""

        reply = await message.reply_media_group(
            media=[
                FileType.media_type(result.file_type)(
                    media=result.photo, caption=result.caption, parse_mode=result.parse_mode, filename=result.filename
                )
                for result in self.results
            ],
            *args,
            **kwargs,
        )

        for index, value in enumerate(reply):
            result = self.results[index]
            await result.cache_file_id(value)