Files
nonebot-plugin-tetris-stats/nonebot_plugin_tetris_stats/utils/screenshot.py

31 lines
1.1 KiB
Python

from playwright.async_api import BrowserContext, TimeoutError, ViewportSize
from ..config.config import config
from .browser import BrowserManager
from .retry import retry
from .time_it import time_it
async def context_factory() -> BrowserContext:
return await (await BrowserManager.get_browser()).new_context(device_scale_factor=config.tetris.screenshot_quality)
@retry(exception_type=TimeoutError, reply='截图失败, 重试中')
@time_it
async def screenshot(url: str) -> bytes:
context = await BrowserManager.get_context('screenshot', factory=context_factory)
async with await context.new_page() as page:
await page.goto(url)
size: ViewportSize = await page.evaluate("""
() => {
const element = document.querySelector('#content');
return {
width: element.offsetWidth,
height: element.offsetHeight,
};
};
""")
await page.set_viewport_size(size)
await page.wait_for_load_state('networkidle')
return await page.locator('id=content').screenshot(animations='disabled', timeout=5000, type='png')