new render

konabot/common/web_render/__init__.py

@@ -1,10 +1,30 @@
 import asyncio
 import queue
 from loguru import logger
-from playwright.async_api import async_playwright, Browser
+from playwright.async_api import async_playwright, Browser, Page, BrowserContext
 
 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: callable = None, timeout: int = 30) -> bytes:
@@ -19,14 +39,46 @@ class WebRenderer:
         :return: 截图的字节数据
 
         '''
-        logger.debug(f"Requesting render for {url} targeting {target} with timeout {timeout}")
-        if cls.browser_pool.empty():
-            instance = await WebRendererInstance.create()
-            cls.browser_pool.put(instance)
-        instance = cls.browser_pool.get()
-        cls.browser_pool.put(instance)
+        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):
@@ -58,28 +110,38 @@ class WebRendererInstance:
         async with self.lock:
             context = await self.browser.new_context()
             page = await context.new_page()
-            logger.debug(f"Navigating to {url} with timeout {timeout}")
-            try:
-                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(f"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:
-                    raise Exception(f"Target element '{target}' not found on the page.")
-                if index >= len(elements):
-                    raise Exception(f"Index {index} out of range for elements matching '{target}'.")
-                element = elements[index]
-                screenshot = await element.screenshot()
-                logger.debug(f"Screenshot taken successfully")
-                return screenshot
-            finally:
-                await page.close()
-                await context.close()
+            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(f"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()