调整 Gif 图渲染策略

konabot/common/nb/extract_image.py

@@ -1,4 +1,5 @@
 from io import BytesIO
+from pathlib import Path
 from typing import Annotated
 
 import httpx
@@ -19,15 +20,21 @@ from PIL import UnidentifiedImageError
 from pydantic import BaseModel
 from returns.result import Failure, Result, Success
 
-from konabot.common.path import ASSETS_PATH
-
 
 discordConfig = nonebot.get_plugin_config(DiscordConfig)
 
 
 class ExtractImageConfig(BaseModel):
     module_extract_image_no_download: bool = False
-    "要不要算了，不下载了，直接爆炸算了，适用于一些比较奇怪的网络环境，无法从协议端下载文件"
+    """
+    要不要算了，不下载了，直接爆炸算了，
+    适用于一些比较奇怪的网络环境，无法从协议端下载文件
+    """
+
+    module_extract_image_target: str = './assets/img/other/boom.jpg'
+    """
+    使用哪个图片呢
+    """
 
 
 module_config = nonebot.get_plugin_config(ExtractImageConfig)
@@ -37,7 +44,7 @@ async def download_image_bytes(url: str, proxy: str | None = None) -> Result[byt
     # if "/matcha/cache/" in url:
     #     url = url.replace('127.0.0.1', '10.126.126.101')
     if module_config.module_extract_image_no_download:
-        return Success((ASSETS_PATH / "img" / "other" / "boom.jpg").read_bytes())
+        return Success(Path(module_config.module_extract_image_target).read_bytes())
     logger.debug(f"开始从 {url} 下载图片")
     async with httpx.AsyncClient(proxy=proxy) as c:
         try:
@@ -70,15 +77,22 @@ def bytes_to_pil(raw_data: bytes | BytesIO) -> Result[PIL.Image.Image, str]:
         return Failure("图像无法读取，可能是网络存在问题orz")
 
 
-async def unimsg_img_to_pil(image: Image) -> Result[PIL.Image.Image, str]:
+async def unimsg_img_to_bytes(image: Image) -> Result[bytes, str]:
     if image.url is not None:
         raw_result = await download_image_bytes(image.url)
     elif image.raw is not None:
-        raw_result = Success(image.raw)
+        if isinstance(image.raw, bytes):
+            raw_result = Success(image.raw)
+        else:
+            raw_result = Success(image.raw.getvalue())
     else:
         return Failure("由于一些内部问题，下载图片失败了orz")
 
-    return raw_result.bind(bytes_to_pil)
+    return raw_result
+
+
+async def unimsg_img_to_pil(image: Image) -> Result[PIL.Image.Image, str]:
+    return (await unimsg_img_to_bytes(image)).bind(bytes_to_pil)
 
 
 async def extract_image_from_qq_message(
@@ -86,7 +100,7 @@ async def extract_image_from_qq_message(
     evt: OnebotV11MessageEvent,
     bot: OnebotV11Bot,
     allow_reply: bool = True,
-) -> Result[PIL.Image.Image, str]:
+) -> Result[bytes, str]:
     if allow_reply and (reply := evt.reply) is not None:
         return await extract_image_from_qq_message(
             reply.message,
@@ -118,18 +132,17 @@ async def extract_image_from_qq_message(
             url = seg.data.get("url")
             if url is None:
                 return Failure("无法下载图片，可能有一些网络问题")
-            data = await download_image_bytes(url)
-            return data.bind(bytes_to_pil)
+            return await download_image_bytes(url)
 
     return Failure("请在消息中包含图片，或者引用一个含有图片的消息")
 
 
-async def extract_image_from_message(
+async def extract_image_data_from_message(
     msg: Message,
     evt: Event,
     bot: Bot,
     allow_reply: bool = True,
-) -> Result[PIL.Image.Image, str]:
+) -> Result[bytes, str]:
     if (
         isinstance(bot, OnebotV11Bot) 
         and isinstance(msg, OnebotV11Message) 
@@ -145,18 +158,18 @@ async def extract_image_from_message(
             if "image/" not in a.content_type:
                 continue
             url = a.proxy_url
-            return (await download_image_bytes(url, discordConfig.discord_proxy)).bind(bytes_to_pil)
+            return await download_image_bytes(url, discordConfig.discord_proxy)
 
     for seg in UniMessage.of(msg, bot):
         logger.info(seg)
         if isinstance(seg, Image):
-            return await unimsg_img_to_pil(seg)
+            return await unimsg_img_to_bytes(seg)
         elif isinstance(seg, Reply) and allow_reply:
             msg2 = seg.msg
             logger.debug(f"深入搜索引用的消息：{msg2}")
             if msg2 is None or isinstance(msg2, str):
                 continue
-            return await extract_image_from_message(msg2, evt, bot, False)
+            return await extract_image_data_from_message(msg2, evt, bot, False)
         elif isinstance(seg, RefNode) and allow_reply:
             if isinstance(bot, DiscordBot):
                 return Failure("暂时不支持在 Discord 中通过引用的方式获取图片")
@@ -165,12 +178,12 @@ async def extract_image_from_message(
     return Failure("请在消息中包含图片，或者引用一个含有图片的消息")
 
 
-async def _ext_img(
+async def _ext_img_data(
     evt: Event,
     bot: Bot,
     matcher: Matcher,
-) -> PIL.Image.Image | None:
-    match await extract_image_from_message(evt.get_message(), evt, bot):
+) -> bytes | None:
+    match await extract_image_data_from_message(evt.get_message(), evt, bot):
         case Success(img):
             return img
         case Failure(err):
@@ -180,4 +193,20 @@ async def _ext_img(
     assert False
 
 
-PIL_Image = Annotated[PIL.Image.Image, nonebot.params.Depends(_ext_img)]
+async def _ext_img(
+    evt: Event,
+    bot: Bot,
+    matcher: Matcher,
+) -> PIL.Image.Image | None:
+    r = await _ext_img_data(evt, bot, matcher)
+    if r:
+        match bytes_to_pil(r):
+            case Success(img):
+                return img
+            case Failure(msg):
+                await matcher.send(await UniMessage.text(msg).export())
+    return None
+
+
+DepImageBytes = Annotated[bytes, nonebot.params.Depends(_ext_img_data)]
+DepPILImage = Annotated[PIL.Image.Image, nonebot.params.Depends(_ext_img)]