PaiGram/core/template/models.py
Chuangbo Li f122e21092
增加 html to image 的缓存
* 增加 html to image 缓存
* 对 template_service.render 进行封装,管理缓存逻辑
* cache key 为 html 的 sha256
* cache value 为 reply_photo 后 telegram 返回的 file_id
* 存入 redis,并设置合理的 ttl

Co-authored-by: 洛水居室 <luoshuijs@outlook.com>
Co-authored-by: xtaodada <xtao@xtaolink.cn>
2022-10-22 15:03:59 +08:00

127 lines
4.1 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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)