184 lines
7.1 KiB
Python
184 lines
7.1 KiB
Python
import asyncio
|
||
import queue
|
||
from typing import Any, Callable, Coroutine
|
||
from loguru import logger
|
||
from playwright.async_api import Page, Playwright, async_playwright, Browser, Page, BrowserContext
|
||
|
||
|
||
PageFunction = Callable[[Page], Coroutine[Any, Any, Any]]
|
||
|
||
|
||
class WebRenderer:
|
||
browser_pool: queue.Queue["WebRendererInstance"] = queue.Queue()
|
||
context_pool: dict[int, BrowserContext] = {} # 长期挂载的浏览器上下文池
|
||
page_pool: dict[str, Page] = {} # 长期挂载的页面池
|
||
|
||
@classmethod
|
||
async def get_browser_instance(cls) -> "WebRendererInstance":
|
||
if cls.browser_pool.empty():
|
||
instance = await WebRendererInstance.create()
|
||
cls.browser_pool.put(instance)
|
||
instance = cls.browser_pool.get()
|
||
cls.browser_pool.put(instance)
|
||
return instance
|
||
|
||
@classmethod
|
||
async def get_browser_context(cls) -> BrowserContext:
|
||
instance = await cls.get_browser_instance()
|
||
if id(instance) not in cls.context_pool:
|
||
context = await instance.browser.new_context()
|
||
cls.context_pool[id(instance)] = context
|
||
logger.debug(f"Created new persistent browser context for WebRendererInstance {id(instance)}")
|
||
return cls.context_pool[id(instance)]
|
||
|
||
@classmethod
|
||
async def render(
|
||
cls,
|
||
url: str,
|
||
target: str,
|
||
params: dict = {},
|
||
other_function: PageFunction | None = None,
|
||
timeout: int = 30,
|
||
) -> bytes:
|
||
'''
|
||
访问指定URL并返回截图
|
||
|
||
:param url: 目标URL
|
||
:param target: 渲染目标，如 ".box"、"#main" 等CSS选择器
|
||
:param timeout: 页面加载超时时间，单位秒
|
||
:param params: URL键值对参数
|
||
:param other_function: 其他自定义操作函数，接受page参数
|
||
:return: 截图的字节数据
|
||
|
||
'''
|
||
instance = await cls.get_browser_instance()
|
||
logger.debug(f"Using WebRendererInstance {id(instance)} to render {url} targeting {target}")
|
||
return await instance.render(url, target, params=params, other_function=other_function, timeout=timeout)
|
||
|
||
|
||
@classmethod
|
||
async def render_persistent_page(cls, page_id: str, url: str, target: str, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes:
|
||
'''
|
||
使用长期挂载的页面访问指定URL并返回截图
|
||
|
||
:param page_id: 页面唯一标识符
|
||
:param url: 目标URL
|
||
:param target: 渲染目标，如 ".box"、"#main" 等CSS选择器
|
||
:param timeout: 页面加载超时时间，单位秒
|
||
:param params: URL键值对参数
|
||
:param other_function: 其他自定义操作函数，接受page参数
|
||
:return: 截图的字节数据
|
||
|
||
'''
|
||
logger.debug(f"Requesting persistent render for page_id {page_id} at {url} targeting {target} with timeout {timeout}")
|
||
instance = await cls.get_browser_instance()
|
||
if page_id not in cls.page_pool:
|
||
context = await cls.get_browser_context()
|
||
page = await context.new_page()
|
||
cls.page_pool[page_id] = page
|
||
logger.debug(f"Created new persistent page for page_id {page_id} using WebRendererInstance {id(instance)}")
|
||
page = cls.page_pool[page_id]
|
||
return await instance.render_with_page(page, url, target, params=params, other_function=other_function, timeout=timeout)
|
||
|
||
@classmethod
|
||
async def close_persistent_page(cls, page_id: str) -> None:
|
||
'''
|
||
关闭并移除长期挂载的页面
|
||
|
||
:param page_id: 页面唯一标识符
|
||
'''
|
||
if page_id in cls.page_pool:
|
||
page = cls.page_pool[page_id]
|
||
await page.close()
|
||
del cls.page_pool[page_id]
|
||
logger.debug(f"Closed and removed persistent page for page_id {page_id}")
|
||
|
||
|
||
|
||
class WebRendererInstance:
|
||
def __init__(self):
|
||
self._playwright: Playwright | None = None
|
||
self._browser: Browser | None = None
|
||
self.lock = asyncio.Lock()
|
||
|
||
@property
|
||
def playwright(self) -> Playwright:
|
||
assert self._playwright is not None
|
||
return self._playwright
|
||
|
||
@property
|
||
def browser(self) -> Browser:
|
||
assert self._browser is not None
|
||
return self._browser
|
||
|
||
async def init(self):
|
||
self._playwright = await async_playwright().start()
|
||
self._browser = await self.playwright.chromium.launch(headless=True)
|
||
|
||
@classmethod
|
||
async def create(cls) -> "WebRendererInstance":
|
||
instance = cls()
|
||
await instance.init()
|
||
return instance
|
||
|
||
async def render(
|
||
self,
|
||
url: str,
|
||
target: str,
|
||
index: int = 0,
|
||
params: dict = {},
|
||
other_function: PageFunction | None = None,
|
||
timeout: int = 30
|
||
) -> bytes:
|
||
'''
|
||
访问指定URL并返回截图
|
||
|
||
:param url: 目标URL
|
||
:param target: 渲染目标，如 ".box"、"#main" 等CSS选择器
|
||
:param timeout: 页面加载超时时间，单位秒
|
||
:param index: 如果目标是一个列表，指定要截图的元素索引
|
||
:param params: URL键值对参数
|
||
:param other_function: 其他自定义操作函数，接受page参数
|
||
:return: 截图的字节数据
|
||
|
||
'''
|
||
async with self.lock:
|
||
context = await self.browser.new_context()
|
||
page = await context.new_page()
|
||
screenshot = await self.inner_render(page, url, target, index, params, other_function, timeout)
|
||
await page.close()
|
||
await context.close()
|
||
return screenshot
|
||
|
||
async def render_with_page(self, page: Page, url: str, target: str, index: int = 0, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes:
|
||
async with self.lock:
|
||
screenshot = await self.inner_render(page, url, target, index, params, other_function, timeout)
|
||
return screenshot
|
||
|
||
async def inner_render(self, page: Page, url: str, target: str, index: int = 0, params: dict = {}, other_function: callable = None, timeout: int = 30) -> bytes:
|
||
logger.debug(f"Navigating to {url} with timeout {timeout}")
|
||
url_with_params = url + ("?" + "&".join(f"{k}={v}" for k, v in params.items()) if params else "")
|
||
await page.goto(url_with_params, timeout=timeout * 1000, wait_until="load")
|
||
logger.debug("Page loaded successfully")
|
||
# 等待目标元素出现
|
||
await page.wait_for_selector(target, timeout=timeout * 1000)
|
||
logger.debug(f"Target element '{target}' found, taking screenshot")
|
||
if other_function:
|
||
await other_function(page)
|
||
elements = await page.query_selector_all(target)
|
||
if not elements:
|
||
logger.error(f"Target element '{target}' not found on the page.")
|
||
return None
|
||
if index >= len(elements):
|
||
logger.error(f"Index {index} out of range for elements matching '{target}'")
|
||
return None
|
||
element = elements[index]
|
||
screenshot = await element.screenshot()
|
||
logger.debug(f"Screenshot taken successfully")
|
||
return screenshot
|
||
|
||
async def close(self):
|
||
await self.browser.close()
|
||
await self.playwright.stop()
|
||
|