🐛 补充主题渲染时遗漏的转发推文中的图片 (#554)

nonebot_bison/theme/themes/basic/build.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
 from nonebot_plugin_saa import Text, Image, MessageSegmentFactory
@@ -44,8 +46,14 @@ class BasicTheme(Theme):
 
         client = await post.platform.ctx.get_client_for_static()
         msgs: list[MessageSegmentFactory] = [Text(text)]
+
+        pics_group: list[list[str | bytes | Path | BytesIO]] = []
         if post.images:
-            pics = post.images
+            pics_group.append(post.images)
+        if rp and rp.images:
+            pics_group.append(rp.images)
+
+        for pics in pics_group:
             if is_pics_mergable(pics):
                 pics = await pic_merge(list(pics), client)
             msgs.extend(map(Image, pics))

nonebot_bison/theme/themes/ceobe_canteen/build.py

@@ -1,3 +1,4 @@
+from io import BytesIO
 from pathlib import Path
 from datetime import datetime
 from typing import TYPE_CHECKING, Literal
@@ -8,6 +9,7 @@ from nonebot_plugin_saa import Text, Image, MessageSegmentFactory
 
 from nonebot_bison.compat import model_validator
 from nonebot_bison.theme.utils import convert_to_qr
+from nonebot_bison.utils.image import pic_merge, is_pics_mergable
 from nonebot_bison.theme import Theme, ThemeRenderError, ThemeRenderUnsupportError
 
 if TYPE_CHECKING:
@@ -109,6 +111,16 @@ class CeobeCanteenTheme(Theme):
             text += f"详情: {post.url}"
         msgs.append(Text(text))
 
+        pics_group: list[list[str | bytes | Path | BytesIO]] = []
         if post.images:
-            msgs.extend(map(Image, post.images))
+            pics_group.append(post.images)
+        if post.repost and post.repost.images:
+            pics_group.append(post.repost.images)
+
+        client = await post.platform.ctx.get_client_for_static()
+        for pics in pics_group:
+            if is_pics_mergable(pics):
+                pics = await pic_merge(list(pics), client)
+            msgs.extend(map(Image, pics))
+
         return msgs

nonebot_bison/theme/themes/ht2i/build.py

@@ -1,3 +1,5 @@
+from io import BytesIO
+from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
 from nonebot_plugin_saa import Text, Image, MessageSegmentFactory
@@ -53,9 +55,15 @@ class Ht2iTheme(Theme):
         if urls:
             msgs.append(Text("\n".join(urls)))
 
+        pics_group: list[list[str | bytes | Path | BytesIO]] = []
         if post.images:
-            client = await post.platform.ctx.get_client_for_static()
+            pics_group.append(post.images)
-            pics = post.images
+        if rp and rp.images:
+            pics_group.append(rp.images)
+
+        client = await post.platform.ctx.get_client_for_static()
+
+        for pics in pics_group:
             if is_pics_mergable(pics):
                 pics = await pic_merge(list(pics), client)
             msgs.extend(map(Image, pics))

tests/platforms/utils.py

@@ -14,3 +14,8 @@ def get_file(file_name: str):
     with open(path / file_name, encoding="utf8") as f:
         file_text = f.read()
         return file_text
+
+
+def get_bytes(file_name: str):
+    with open(path / file_name, "rb") as f:
+        return f.read()

tests/theme/test_themes.py

@@ -1,11 +1,15 @@
 from time import time
-from typing import Any
+from copy import deepcopy
 from inspect import cleandoc
+from typing import TYPE_CHECKING, Any
 
 import pytest
 from flaky import flaky
 from nonebug import App
 
+if TYPE_CHECKING:
+    from nonebot_bison.post import Post
+
 now = time()
 passed = now - 3 * 60 * 60
 raw_post_list_1 = [{"id": 1, "text": "p1", "date": now, "tags": ["tag1"], "category": 1}]
@@ -17,6 +21,18 @@ raw_post_list_2 = raw_post_list_1 + [
 ]
 
 
+@pytest.fixture()
+def MERGEABLE_PNG_DATA() -> bytes:
+    from tests.platforms.utils import get_bytes
+
+    return get_bytes("mergeable-pic.jpg")
+
+
+@pytest.fixture()
+def SIMPLE_PNG_DATA() -> bytes:
+    return b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89"
+
+
 @pytest.fixture()
 def mock_platform(app: App):
     from nonebot_bison.post import Post
@@ -160,7 +176,7 @@ async def test_arknights_theme(app: App, mock_post):
 
 
 @pytest.mark.asyncio
-async def test_basic_theme(app: App, mock_post):
+async def test_basic_theme(app: App, mock_post: "Post", MERGEABLE_PNG_DATA, SIMPLE_PNG_DATA):
     from nonebot_plugin_saa import Text, Image
 
     from nonebot_bison.theme import theme_manager
@@ -171,7 +187,7 @@ async def test_basic_theme(app: App, mock_post):
     assert isinstance(basic_theme, BasicTheme)
     assert basic_theme.name == "basic"
     res = await basic_theme.render(mock_post)
-    assert len(res) == 2
+    assert len(res) == 3
     assert res[0] == Text(
         cleandoc(
             """title
@@ -190,7 +206,7 @@ async def test_basic_theme(app: App, mock_post):
     )
     assert isinstance(res[1], Image)
 
-    mock_post2 = mock_post
+    mock_post2 = deepcopy(mock_post)
     mock_post2.repost = None
     mock_post2.images = []
     res2 = await basic_theme.render(mock_post2)
@@ -206,6 +222,13 @@ async def test_basic_theme(app: App, mock_post):
         )
     )
 
+    mock_post3 = deepcopy(mock_post)
+    mock_post3.images = [MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, SIMPLE_PNG_DATA]
+    assert mock_post3.repost
+    mock_post3.repost.images = [SIMPLE_PNG_DATA, SIMPLE_PNG_DATA]
+    res3 = await basic_theme.render(mock_post3)
+    assert len(res3) == 5
+
 
 @pytest.mark.asyncio
 async def test_brief_theme(app: App, mock_post):
@@ -250,7 +273,7 @@ async def test_brief_theme(app: App, mock_post):
 @pytest.mark.render
 @pytest.mark.asyncio
 @flaky(max_runs=3, min_passes=1)
-async def test_ceobecanteen_theme(app: App, mock_post):
+async def test_ceobecanteen_theme(app: App, mock_post: "Post", MERGEABLE_PNG_DATA, SIMPLE_PNG_DATA):
     from nonebot_plugin_saa import Text, Image
 
     from nonebot_bison.theme import theme_manager
@@ -261,16 +284,22 @@ async def test_ceobecanteen_theme(app: App, mock_post):
     assert isinstance(ceobecanteen_theme, CeobeCanteenTheme)
     assert ceobecanteen_theme.name == "ceobecanteen"
     res = await ceobecanteen_theme.render(mock_post)
-    assert len(res) == 3
+    assert len(res) == 4
     assert isinstance(res[0], Image)
     assert isinstance(res[2], Image)
     assert res[1] == Text("来源: Mock Platform Mock\n详情: http://t.tt/1")
 
+    mock_post2 = deepcopy(mock_post)
+    assert mock_post2.repost
+    mock_post2.repost.images = [MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, SIMPLE_PNG_DATA]
+    res2 = await ceobecanteen_theme.render(mock_post2)
+    assert len(res2) == 5
+
 
 @pytest.mark.render
 @pytest.mark.asyncio
 @flaky(max_runs=3, min_passes=1)
-async def test_ht2i_theme(app: App, mock_post):
+async def test_ht2i_theme(app: App, mock_post: "Post", MERGEABLE_PNG_DATA, SIMPLE_PNG_DATA):
     from nonebot_plugin_saa import Text, Image
 
     from nonebot_bison.theme import theme_manager
@@ -281,16 +310,16 @@ async def test_ht2i_theme(app: App, mock_post):
     assert isinstance(ht2i_theme, Ht2iTheme)
     assert ht2i_theme.name == "ht2i"
     res = await ht2i_theme.render(mock_post)
-    assert len(res) == 3
+    assert len(res) == 4
     assert isinstance(res[0], Image)
     assert isinstance(res[2], Image)
     assert res[1] == Text("转发详情: http://t.tt/2\n详情: http://t.tt/1")
 
-    mock_post2 = mock_post
+    mock_post2 = deepcopy(mock_post)
     mock_post2.repost = None
-    mock_post2.images = []
+    mock_post2.images = [MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA, MERGEABLE_PNG_DATA]
 
     res2 = await ht2i_theme.render(mock_post2)
 
-    assert len(res2) == 2
+    assert len(res2) == 4
     assert res2[1] == Text("详情: http://t.tt/1")