update pve.py

nonebot_plugin_fortnite/pve.py

@@ -1,9 +1,7 @@
 import asyncio
-
 from PIL import Image
 from pathlib import Path
 from playwright.async_api import async_playwright
-
 from .config import data_dir
 
 vb_file = data_dir / "vb.png"
@@ -16,53 +14,55 @@ async def screenshot_vb_img() -> Path:
         async with async_playwright() as p:
             browser = await p.chromium.launch(headless=True)
             page = await browser.new_page()
-            await page.goto(url)  
-            # 截取第一个 <div class="hot-info">
+            await page.goto(url)
+
+            # Check and screenshot the first <div class="hot-info">
             hot_info_1 = page.locator('div.hot-info').nth(0)
-            await hot_info_1.screenshot(path=data_dir / 'hot_info_1.png')
-
-            # 截取 <div class="container hidden-xs">
+            if await hot_info_1.count() > 0:
+                await hot_info_1.screenshot(path=data_dir / 'hot_info_1.png')
+
+            # Check and screenshot <div class="container hidden-xs">
             container_hidden_xs = page.locator('div.container.hidden-xs')
-            await container_hidden_xs.screenshot(path=data_dir / 'container_hidden_xs.png')
-            # 截取第二个 <div class="hot-info">
+            if await container_hidden_xs.count() > 0:
+                await container_hidden_xs.screenshot(path=data_dir / 'container_hidden_xs.png')
+
+            # Check and screenshot the second <div class="hot-info">
             hot_info_2 = page.locator('div.hot-info').nth(1)
-            await hot_info_2.screenshot(path=data_dir / 'hot_info_2.png')
+            if await hot_info_2.count() > 0:
+                await hot_info_2.screenshot(path=data_dir / 'hot_info_2.png')
+
             combine_imgs()
     finally:
         if browser:
             await browser.close()
 
     return vb_file
 
-
-
 def combine_imgs():
     try:
-        # 打开三个截图
-        hot_info_1 = Image.open(data_dir / 'hot_info_1.png')
-        container_hidden_xs = Image.open(data_dir / 'container_hidden_xs.png')
-        hot_info_2 = Image.open(data_dir / 'hot_info_2.png')
-
-        # 获取每个截图的尺寸
-        width1, height1 = hot_info_1.size
-        width2, height2 = container_hidden_xs.size
-        width3, height3 = hot_info_2.size
+        # Open the screenshots if they exist
+        images = []
+        for img_name in ['hot_info_1.png', 'container_hidden_xs.png', 'hot_info_2.png']:
+            img_path = data_dir / img_name
+            if img_path.exists():
+                images.append(Image.open(img_path))
 
-        # 创建一个新的图像，宽度为最宽的截图宽度，高度为三个截图高度之和
-        total_width = max(width1, width2, width3)
-        total_height = height1 + height2 + height3
+        # Get dimensions and create a new image
+        widths, heights = zip(*(img.size for img in images))
+        total_width = max(widths)
+        total_height = sum(heights)
         combined_image = Image.new('RGB', (total_width, total_height))
 
-        # 将每个截图粘贴到新的图像中
-        combined_image.paste(hot_info_1, (0, 0))
-        combined_image.paste(container_hidden_xs, (0, height1))
-        combined_image.paste(hot_info_2, (0, height1 + height2))
+        # Paste images into the new image
+        y_offset = 0
+        for img in images:
+            combined_image.paste(img, (0, y_offset))
+            y_offset += img.height
 
-        # 保存合并后的图像
+        # Save the combined image
         combined_image.save(vb_file)
     finally:
-        # 关闭所有打开的图像资源
-        hot_info_1.close()
-        container_hidden_xs.close()
-        hot_info_2.close()
+        # Close all opened images
+        for img in images:
+            img.close()
         combined_image.close()