feat: 添加了支持选择元素截图和等待资源加载完成之后再截图的功能

src/render.py

@@ -1,4 +1,6 @@
+import asyncio
 import re
+import time
 
 from .util import generate_data_path
 from playwright.async_api import async_playwright
@@ -8,6 +10,8 @@
 from typing import Literal
 from loguru import logger
 from playwright.async_api import BrowserContext, Browser, Playwright
+from playwright.async_api import Error as PlaywrightError
+from playwright.async_api import TimeoutError as PlaywrightTimeoutError
 from playwright._impl._errors import TargetClosedError
 
 
@@ -51,6 +55,13 @@ class ScreenshotOptions(BaseModel):
             - normal: 1.0
             - high: 1.3
             - ultra: 1.8
+        selector: (str, optional): CSS 选择器。设置后优先对匹配的元素截图，而不是整页截图.
+        fallback_selector: (str, optional): selector 未命中时使用的备用 CSS 选择器.
+        selector_timeout: (int, optional): 等待 selector 出现的超时时间，单位毫秒.
+        wait_until: (Literal["commit", "domcontentloaded", "load", "networkidle"], optional):
+            页面导航等待状态。未指定时使用 Playwright 默认值.
+        wait_for_resources: (bool, optional): 是否在截图前等待网络空闲、图片解码和字体加载.
+        resource_timeout: (int, optional): 等待资源加载的超时时间，单位毫秒，默认 5000.
 
     @author: Redlnn(https://github.com/GraiaCommunity/graiax-text2img-playwright)
     """
@@ -67,6 +78,12 @@ class ScreenshotOptions(BaseModel):
     viewport_width: int | None = None
     viewport_height: int | None = None
     device_scale_factor_level: Literal["normal", "high", "ultra", None] = None
+    selector: str | None = None
+    fallback_selector: str | None = None
+    selector_timeout: int | None = None
+    wait_until: Literal["commit", "domcontentloaded", "load", "networkidle", None] = None
+    wait_for_resources: bool | None = None
+    resource_timeout: int | None = None
 
 
 class Text2ImgRender:
@@ -242,23 +259,197 @@ async def html2pic(
             logger.info(f"html2pic: set viewport size to {width}x{height}")
 
         try:
-            await page.goto(
-                f"file://{html_file_path}", timeout=screenshot_options.timeout
-            )
+            goto_kwargs = {"timeout": screenshot_options.timeout}
+            if screenshot_options.wait_until:
+                goto_kwargs["wait_until"] = screenshot_options.wait_until
+            await page.goto(f"file://{html_file_path}", **goto_kwargs)
+
+            if screenshot_options.wait_for_resources:
+                await self._wait_for_resources(page, screenshot_options)
+
             screenshot_kwargs = screenshot_options.model_dump(exclude_none=True)
             screenshot_kwargs.pop("viewport_width", None)
             screenshot_kwargs.pop("viewport_height", None)
             screenshot_kwargs.pop("device_scale_factor_level", None)
+            screenshot_kwargs.pop("selector", None)
+            screenshot_kwargs.pop("fallback_selector", None)
+            screenshot_kwargs.pop("selector_timeout", None)
+            screenshot_kwargs.pop("wait_until", None)
+            screenshot_kwargs.pop("wait_for_resources", None)
+            screenshot_kwargs.pop("resource_timeout", None)
 
             # Robustness: Remove quality if type is png, as Playwright errors out
             if screenshot_options.type == "png":
                 screenshot_kwargs.pop("quality", None)
 
-            await page.screenshot(path=result_path, **screenshot_kwargs)
+            element = await self._resolve_screenshot_element(page, screenshot_options)
+            if element is not None:
+                element_screenshot_kwargs = screenshot_kwargs.copy()
+                element_screenshot_kwargs.pop("full_page", None)
+                element_screenshot_kwargs.pop("clip", None)
+                await element.screenshot(path=result_path, **element_screenshot_kwargs)
+            else:
+                await page.screenshot(path=result_path, **screenshot_kwargs)
         finally:
             # Ensure the page is closed to free resources
             await page.close()
 
         logger.info(f"Rendered {html_file_path} to {result_path}")
 
         return result_path
+
+    async def _resolve_screenshot_element(self, page, screenshot_options: ScreenshotOptions):
+        """Resolve an element target for screenshot when selector options are provided."""
+
+        if not screenshot_options.selector:
+            return None
+
+        element = None
+        if screenshot_options.selector_timeout is not None:
+            try:
+                element = await page.wait_for_selector(
+                    screenshot_options.selector,
+                    timeout=screenshot_options.selector_timeout,
+                )
+            except PlaywrightTimeoutError as e:
+                logger.debug(
+                    f"html2pic: wait for selector '{screenshot_options.selector}' failed: {e}"
+                )
+        else:
+            try:
+                element = await page.query_selector(screenshot_options.selector)
+            except PlaywrightError as e:
+                logger.warning(
+                    f"html2pic: invalid selector '{screenshot_options.selector}', "
+                    f"skip element screenshot: {e}"
+                )
+
+        if element is not None:
+            logger.info(
+                f"html2pic: screenshot element matched selector '{screenshot_options.selector}'"
+            )
+            return element
+
+        if screenshot_options.fallback_selector:
+            try:
+                fallback = await page.query_selector(screenshot_options.fallback_selector)
+            except PlaywrightError as e:
+                logger.warning(
+                    "html2pic: invalid fallback selector "
+                    f"'{screenshot_options.fallback_selector}', fallback to page screenshot: {e}"
+                )
+                fallback = None
+            if fallback is not None:
+                logger.info(
+                    f"html2pic: selector '{screenshot_options.selector}' not found, "
+                    f"screenshot fallback selector '{screenshot_options.fallback_selector}'"
+                )
+                return fallback
+
+        logger.warning(
+            f"html2pic: selector '{screenshot_options.selector}' not found, fallback to page screenshot"
+        )
+        return None
+
+    async def _wait_for_resources(self, page, screenshot_options: ScreenshotOptions):
+        """Wait for common remote resources before taking a screenshot."""
+
+        timeout = screenshot_options.resource_timeout
+        if timeout is None:
+            timeout = screenshot_options.timeout if screenshot_options.timeout is not None else 5000
+        if timeout <= 0:
+            return
+
+        deadline = time.monotonic() + timeout / 1000
+
+        def remaining_timeout() -> int:
+            return max(0, int((deadline - time.monotonic()) * 1000))
+
+        try:
+            network_timeout = remaining_timeout()
+            if network_timeout > 0:
+                await page.wait_for_load_state("networkidle", timeout=network_timeout)
+        except Exception as e:
+            logger.warning(f"html2pic: wait for networkidle failed: {e}")
+
+        resource_timeout = remaining_timeout()
+        if resource_timeout <= 0:
+            logger.warning(f"html2pic: wait for resources timed out after {timeout}ms")
+            return
+
+        try:
+            eval_timeout = resource_timeout
+            if screenshot_options.timeout is not None and screenshot_options.timeout > 0:
+                eval_timeout = min(resource_timeout, int(screenshot_options.timeout))
+
+            result = await asyncio.wait_for(
+                page.evaluate(
+                    """
+                async (timeout) => {
+                  const images = Array.from(document.images || []);
+                  const waitImage = (img) => {
+                    if (img.complete) {
+                      return Promise.resolve();
+                    }
+                    if (typeof img.decode === "function") {
+                      return img.decode().catch(() => undefined);
+                    }
+                    return new Promise((resolve) => {
+                      img.addEventListener("load", resolve, { once: true });
+                      img.addEventListener("error", resolve, { once: true });
+                    });
+                  };
+                  const waitFonts = () => {
+                    if (document.fonts && document.fonts.ready) {
+                      return document.fonts.ready.catch(() => undefined);
+                    }
+                    return Promise.resolve();
+                  };
+                  let timedOut = false;
+                  const allResources = Promise.all([
+                    ...images.map(waitImage),
+                    waitFonts(),
+                  ]);
+                  const timeoutPromise = new Promise((resolve) => {
+                    setTimeout(() => {
+                      timedOut = true;
+                      resolve();
+                    }, timeout);
+                  });
+                  await Promise.race([allResources, timeoutPromise]);
+                  return {
+                    imageCount: images.length,
+                    completeCount: images.filter((img) => img.complete).length,
+                    brokenCount: images.filter((img) => img.complete && img.naturalWidth === 0).length,
+                    timedOut,
+                  };
+                }
+                """,
+                    eval_timeout,
+                ),
+                timeout=eval_timeout / 1000,
+            )
+            if not isinstance(result, dict):
+                logger.warning("html2pic: wait for resources returned invalid result")
+                return
+
+            if result.get("timedOut"):
+                logger.warning(
+                    "html2pic: wait for resources timed out after "
+                    f"{timeout}ms, images={result.get('completeCount')}/"
+                    f"{result.get('imageCount')}"
+                )
+            elif result.get("brokenCount"):
+                logger.warning(
+                    "html2pic: resources loaded with broken images, "
+                    f"broken={result.get('brokenCount')}/"
+                    f"{result.get('imageCount')}"
+                )
+            else:
+                logger.info(
+                    "html2pic: resources ready, "
+                    f"images={result.get('completeCount')}/"
+                    f"{result.get('imageCount')}"
+                )
+        except Exception as e:
+            logger.warning(f"html2pic: wait for image/font resources failed: {e}")

update.md

@@ -0,0 +1,103 @@
+# T2I 渲染配置更新说明
+
+本次新增两组可选配置：
+
+- `selector screenshot`：按指定 HTML 元素截图，避免整页截图出现多余边缘。
+- `resource wait`：在截图前等待远程图片等资源加载完成。
+
+## Selector Screenshot 配置项
+
+| 配置项 | 类型 | 不配置时的默认值 | 说明 |
+| --- | --- | --- | --- |
+| `selector` | `str` | `None` | CSS 选择器。配置后优先对匹配到的元素执行截图，而不是对整个页面截图。 |
+| `fallback_selector` | `str` | `None` | `selector` 未匹配到元素时使用的备用 CSS 选择器。建议配置为 `body`。 |
+| `selector_timeout` | `int` | `None` | 等待 `selector` 出现的超时时间，单位为毫秒。未配置时不会额外等待，只会立即查询当前页面中是否存在该元素。 |
+
+### 不配置时的默认行为
+
+如果不配置 `selector`，服务端保持原来的行为，继续使用 `page.screenshot()` 对页面截图。
+
+如果配置了 `selector`，但没有配置 `fallback_selector`，并且 `selector` 没有匹配到元素，服务端会记录 warning，然后回退到原来的页面截图行为。
+
+### JSON 请求示例
+
+```json
+{
+  "tmpl": "<div id=\"container\">Hello</div>",
+  "tmpldata": {},
+  "json": true,
+  "options": {
+    "selector": "#container",
+    "fallback_selector": "body",
+    "selector_timeout": 1000
+  }
+}
+```
+
+### AstrBot 插件调用示例
+
+```python
+await self.html_render(
+    tmpl,
+    data,
+    return_url=False,
+    options={
+        "selector": "#container",
+        "fallback_selector": "body",
+        "selector_timeout": 1000,
+    },
+)
+```
+
+## Resource Wait 配置项
+
+| 配置项 | 类型 | 不配置时的默认值 | 说明 |
+| --- | --- | --- | --- |
+| `wait_for_resources` | `bool` | `None`，等效于 `False` | 是否启用截图前资源等待。未配置时不启用额外等待，保持原来的渲染行为。 |
+| `resource_timeout` | `int` | `None` | 资源等待总超时时间，单位为毫秒。仅在 `wait_for_resources` 为 `true` 时生效。未配置时优先使用已有的 `timeout`；如果 `timeout` 也未配置，则默认使用 `5000` 毫秒。 |
+
+### 配置方法
+
+这两个配置项放在 t2i 请求体的 `options` 字段中，不放在 `tmpldata` 中。
+
+### JSON 请求示例
+
+```json
+{
+  "tmpl": "<div id=\"container\"><img src=\"{{ image_url }}\"></div>",
+  "tmpldata": {
+    "image_url": "https://example.com/image.png"
+  },
+  "json": true,
+  "options": {
+    "wait_for_resources": true,
+    "resource_timeout": 10000
+  }
+}
+```
+
+### AstrBot 插件调用示例
+
+```python
+await self.html_render(
+    tmpl,
+    data,
+    return_url=False,
+    options={
+        "wait_for_resources": True,
+        "resource_timeout": 10000,
+    },
+)
+```
+
+### 行为说明
+
+启用 `wait_for_resources` 后，服务端会在截图前等待以下资源状态：
+
+- Playwright `networkidle`
+- 页面内所有 `<img>` 元素加载或解码完成
+- `document.fonts.ready`
+
+如果资源很快加载完成，会立即继续截图，不会固定等待到 `resource_timeout`。
+
+如果超过 `resource_timeout` 仍未加载完成，会记录 warning 并继续截图，避免单个慢速或失效资源导致整个渲染请求永久阻塞。