nonebot-plugin-htmlrender 0.6.1__tar.gz → 0.6.3__tar.gz

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nonebot-plugin-htmlrender
-Version: 0.6.1
+Version: 0.6.3
 Summary: 通过浏览器渲染图片
 Project-URL: Homepage, https://github.com/kexue-z/nonebot-plugin-htmlrender
 Project-URL: Bug Tracker, https://github.com/kexue-z/nonebot-plugin-htmlrender/issues

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/nonebot_plugin_htmlrender/browser.py

@@ -53,15 +53,13 @@ async def init_browser(**kwargs) -> Browser:
     Raises:
         RuntimeError: 如果浏览器无法启动或安装失败。
     """
-    await install_browser()  # update playwright when start
-    await check_playwright_env()
     return await start_browser(**kwargs)
 
 
 @asynccontextmanager
 async def get_new_page(device_scale_factor: float = 2, **kwargs) -> AsyncIterator[Page]:
     """
-    获取一个新的页面的上下文管理器。
+    获取一个新的页面的上下文管理器, 这里的 page 默认使用设备缩放因子为 2。
 
     Args:
         device_scale_factor (float): 设备缩放因子。
@@ -117,6 +115,31 @@ async def _connect_via_cdp(**kwargs) -> Browser:
         raise RuntimeError("Playwright 未初始化")
 
 
+async def _connect(browser_type: str, **kwargs) -> Browser:
+    """
+    通过 Playwright 协议连接浏览器。
+
+    Args:
+        browser_type (str): 浏览器类型。
+        **kwargs: 传递给`playwright.connect`的关键字参数。
+
+    Returns:
+        Browser: 启动的浏览器实例。
+
+    Raises:
+        RuntimeError: 如果 Playwright 未初始化。
+    """
+    _browser_cls: BrowserType = getattr(_playwright, browser_type)
+    kwargs["ws_endpoint"] = plugin_config.htmlrender_connect
+    logger.info(
+        f"正在使用 Playwright 协议连接 {browser_type}({plugin_config.htmlrender_connect})"
+    )
+    if _playwright is not None:
+        return await _browser_cls.connect(**kwargs)
+    else:
+        raise RuntimeError("Playwright 未初始化")
+
+
 async def start_browser(**kwargs) -> Browser:
     """
     启动 Playwright 浏览器实例。
@@ -136,6 +159,9 @@ async def start_browser(**kwargs) -> Browser:
     ):
         return await _connect_via_cdp(**kwargs)
 
+    if plugin_config.htmlrender_connect:
+        return await _connect(plugin_config.htmlrender_browser, **kwargs)
+
     if plugin_config.htmlrender_browser_channel:
         kwargs["channel"] = plugin_config.htmlrender_browser_channel
 
@@ -144,6 +170,12 @@ async def start_browser(**kwargs) -> Browser:
 
     if plugin_config.htmlrender_browser_executable_path:
         kwargs["executable_path"] = plugin_config.htmlrender_browser_executable_path
+    else:
+        try:
+            await check_playwright_env()
+        except RuntimeError:
+            await install_browser()
+            await check_playwright_env()
 
     _browser = await _launch(plugin_config.htmlrender_browser, **kwargs)
     return _browser
@@ -164,7 +196,7 @@ async def check_playwright_env():
     logger.info("Checking Playwright environment...")
     try:
         async with async_playwright() as p:
-            await p.chromium.launch()
+            await getattr(p, plugin_config.htmlrender_browser).launch()
     except Exception as e:
         raise RuntimeError(
             "Playwright environment is not set up correctly. "

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/nonebot_plugin_htmlrender/config.py

@@ -36,6 +36,9 @@ class Config(BaseModel):
     htmlrender_connect_over_cdp: Optional[str] = Field(
         default=None, description="通过 CDP 连接Playwright浏览器的端点地址。"
     )
+    htmlrender_connect: Optional[str] = Field(
+        default=None, description="通过Playwright协议连接Playwright浏览器的端点地址。"
+    )
 
     @model_validator(mode="after")
     def check_browser_channel(self) -> Self:

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/nonebot_plugin_htmlrender/data_source.py

@@ -284,17 +284,32 @@ async def template_to_pic(
 async def capture_element(
     url: str,
     element: str,
-    timeout: float = 0,
-    type: Literal["jpeg", "png"] = "png",
-    quality: Union[int, None] = None,
-    screenshot_timeout: Optional[float] = 30_000,
-    **kwargs,
+    page_kwargs: Optional[dict] = None,
+    goto_kwargs: Optional[dict] = None,
+    screenshot_kwargs: Optional[dict] = None,
 ) -> bytes:
-    async with get_new_page(**kwargs) as page:
-        page.on("console", lambda msg: logger.debug(f"浏览器控制台: {msg.text}"))
-        await page.goto(url, timeout=timeout)
-        return await page.locator(element).screenshot(
-            type=type,
-            quality=quality,
-            timeout=screenshot_timeout,
+    """捕获网页中指定元素的截图, 通过CSS选择器或XPath表达式指定元素。
+
+    Args:
+        url: 目标网页URL
+        element: CSS选择器或XPath表达式
+        page_kwargs: 传递给get_new_page的参数
+        goto_kwargs: 传递给page.goto方法的额外参数
+        screenshot_kwargs: 传递给screenshot方法的额外参数
+
+    Returns:
+        bytes: 元素截图数据
+    """
+    page_kwargs = page_kwargs or {}
+    goto_kwargs = goto_kwargs or {}
+    screenshot_kwargs = screenshot_kwargs or {}
+
+    async with get_new_page(**page_kwargs) as page:
+        page.on(
+            "console",
+            lambda msg: logger.opt(colors=True).debug(
+                f"<cyan>[Browser Console]</cyan> {msg.text}"
+            ),
         )
+        await page.goto(url, **goto_kwargs)
+        return await page.locator(element).screenshot(**screenshot_kwargs)

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nonebot-plugin-htmlrender"
-version = "0.6.1"
+version = "0.6.3"
 description = "通过浏览器渲染图片"
 readme = "README.md"
 authors = [{ name = "kexue", email = "x@kexue-cloud.cn" }]

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/tests/test_browser.py

@@ -39,6 +39,7 @@ def browser_config() -> dict[str, str]:
     return {
         "browser": "chromium",
         "cdp": "ws://localhost:9222",
+        "pwp": "ws://localhost:3000/chromium/playwright",
     }
 
 
@@ -169,6 +170,29 @@ async def test_connect_via_cdp(
     assert browser == mock_browser
 
 
+@pytest.mark.asyncio
+async def test_connect(
+    mocker: MockerFixture, mock_browser: Browser, browser_config: dict[str, str]
+) -> None:
+    """测试通过Playwright协议连接浏览器"""
+    from nonebot_plugin_htmlrender.browser import start_browser
+
+    mocker.patch(
+        "nonebot_plugin_htmlrender.browser._connect", return_value=mock_browser
+    )
+    mocker.patch(
+        "nonebot_plugin_htmlrender.browser.plugin_config.htmlrender_browser",
+        browser_config["browser"],
+    )
+    mocker.patch(
+        "nonebot_plugin_htmlrender.browser.plugin_config.htmlrender_connect",
+        browser_config["pwp"],
+    )
+
+    browser = await start_browser()
+    assert browser == mock_browser
+
+
 @pytest.mark.parametrize(
     ("proxy_url", "expected"),
     [

{nonebot_plugin_htmlrender-0.6.1 → nonebot_plugin_htmlrender-0.6.3}/tests/test_htmlrender.py

@@ -5,6 +5,7 @@ from typing import Any
 from nonebug import App
 from PIL import Image, ImageChops
 import pytest
+from pytest_mock import MockerFixture
 
 
 @pytest.fixture
@@ -39,7 +40,7 @@ def template_resources(request: Any) -> tuple[str, str, list[str]]:
     elif template_type == "text":
         template_name = "text.html"
         data_list = ["1", "2", "3", "4"]
-    else:
+    else:  # pragma: no cover
         raise ValueError(f"Unsupported template type: {template_type}")
 
     return template_path, template_name, data_list
@@ -137,3 +138,75 @@ async def test_template_filter(
     image = Image.open(BytesIO(initial_bytes=image_byte))
     diff = ImageChops.difference(image, test_image)
     assert diff.getbbox() is None
+
+
+@pytest.mark.asyncio
+async def test_capture_element(mocker: MockerFixture) -> None:
+    """测试网页元素捕获功能"""
+    from nonebot_plugin_htmlrender.data_source import capture_element
+
+    mock_screenshot = b"test_image_bytes"
+
+    mock_locator = mocker.AsyncMock()
+    mock_locator.screenshot.return_value = mock_screenshot
+
+    mock_page = mocker.AsyncMock()
+    mock_page.goto = mocker.AsyncMock()
+    mock_page.on = mocker.AsyncMock()
+    mock_page.locator = mocker.MagicMock(return_value=mock_locator)
+
+    mock_cm = mocker.MagicMock()
+    mock_cm.__aenter__ = mocker.AsyncMock(return_value=mock_page)
+    mock_cm.__aexit__ = mocker.AsyncMock(return_value=None)
+
+    mocker.patch(
+        "nonebot_plugin_htmlrender.data_source.get_new_page", return_value=mock_cm
+    )
+
+    result = await capture_element("https://example.com", "#target-element")
+
+    assert result == mock_screenshot
+    mock_page.goto.assert_called_once_with("https://example.com")
+    mock_page.locator.assert_called_once_with("#target-element")
+    mock_locator.screenshot.assert_called_once_with()
+
+    mock_page.goto.reset_mock()
+    mock_page.locator.reset_mock()
+    mock_locator.screenshot.reset_mock()
+
+    page_kwargs = {"device_scale_factor": 2.0}
+    goto_kwargs = {"timeout": 5000}
+    screenshot_kwargs = {"type": "jpeg", "quality": 80}
+
+    result = await capture_element(
+        "https://example.com",
+        "//div[@id='xpath-element']",
+        page_kwargs=page_kwargs,
+        goto_kwargs=goto_kwargs,
+        screenshot_kwargs=screenshot_kwargs,
+    )
+
+    assert result == mock_screenshot
+    mock_page.goto.assert_called_once_with("https://example.com", timeout=5000)
+    mock_page.locator.assert_called_once_with("//div[@id='xpath-element']")
+    mock_locator.screenshot.assert_called_once_with(type="jpeg", quality=80)
+
+
+@pytest.mark.asyncio
+async def test_capture_element_exceptions_propagate(mocker: MockerFixture) -> None:
+    """测试网页元素捕获时的异常能正确传递"""
+    from playwright.async_api import Error
+
+    from nonebot_plugin_htmlrender.data_source import capture_element
+
+    mock_cm = mocker.MagicMock()
+    mock_cm.__aenter__ = mocker.AsyncMock(side_effect=Error("Browser error"))
+
+    mocker.patch(
+        "nonebot_plugin_htmlrender.data_source.get_new_page", return_value=mock_cm
+    )
+
+    with pytest.raises(Error) as exc_info:
+        await capture_element("https://example.com", "#element")
+
+    assert "Browser error" in str(exc_info.value)