K2cr2O1
caf5b06097
feat(ws_pool): 新增 WebSocket 连接池实现 perf(json): 使用 orjson 替代标准 json 库提升性能 style: 清理未使用的导入和冗余代码 docs: 更新架构文档和开发规范 test: 添加 WebSocket 连接池测试用例 fix(plugins): 修复自动审批插件 API 调用参数格式
132 lines
5.0 KiB
Python
132 lines
5.0 KiB
Python
"""
|
||
图片生成管理器模块
|
||
|
||
负责管理图片生成相关的逻辑，支持多种渲染引擎（目前支持 Playwright）。
|
||
"""
|
||
import os
|
||
import base64
|
||
from typing import Dict, Any, Optional
|
||
from jinja2 import Template
|
||
|
||
from .browser_manager import browser_manager
|
||
from ..utils.logger import logger
|
||
from ..utils.singleton import Singleton
|
||
|
||
class ImageManager(Singleton):
|
||
"""
|
||
图片生成管理器（单例）
|
||
"""
|
||
|
||
def __init__(self):
|
||
"""
|
||
初始化图片生成管理器
|
||
"""
|
||
# 检查是否已经初始化
|
||
if hasattr(self, 'template_dir'):
|
||
return
|
||
|
||
# 模板目录
|
||
self.template_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "templates")
|
||
# 临时文件目录
|
||
# core/managers/image_manager.py -> core/managers -> core -> core/data/temp
|
||
self.temp_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "temp")
|
||
os.makedirs(self.temp_dir, exist_ok=True)
|
||
# 模板缓存
|
||
self._template_cache: Dict[str, Template] = {}
|
||
|
||
async def render_template(self, template_name: str, data: Dict[str, Any], output_name: str = "output.png", quality: int = 80, image_type: str = "png") -> Optional[str]:
|
||
"""
|
||
使用 Playwright 渲染 Jinja2 模板并保存为图片文件
|
||
|
||
Args:
|
||
template_name (str): 模板文件名 (例如 "help.html")
|
||
data (Dict[str, Any]): 传递给模板的数据字典
|
||
output_name (str, optional): 输出文件名. Defaults to "output.png".
|
||
quality (int, optional): JPEG 质量 (0-100). 仅在 image_type 为 jpeg 时有效. Defaults to 80.
|
||
image_type (str, optional): 图片类型 ('png' or 'jpeg'). Defaults to "png".
|
||
|
||
Returns:
|
||
Optional[str]: 生成图片的绝对路径，如果失败则返回 None
|
||
"""
|
||
template_path = os.path.join(self.template_dir, template_name)
|
||
if not os.path.exists(template_path):
|
||
logger.error(f"模板文件未找到: {template_path}")
|
||
return None
|
||
|
||
try:
|
||
# 1. 渲染 HTML (使用缓存)
|
||
if template_name in self._template_cache:
|
||
template = self._template_cache[template_name]
|
||
else:
|
||
with open(template_path, "r", encoding="utf-8") as f:
|
||
template_str = f.read()
|
||
template = Template(template_str)
|
||
self._template_cache[template_name] = template
|
||
|
||
html_content = template.render(**data)
|
||
|
||
# 2. 使用浏览器截图
|
||
# 改为从池中获取页面
|
||
page = await browser_manager.get_page()
|
||
if not page:
|
||
logger.error("无法获取浏览器页面")
|
||
return None
|
||
|
||
try:
|
||
width = 1920
|
||
height = 1080
|
||
await page.set_viewport_size({"width": width, "height": height})
|
||
|
||
# 加载内容
|
||
await page.set_content(html_content)
|
||
await page.wait_for_selector("body")
|
||
|
||
|
||
screenshot_args = {
|
||
'full_page': True,
|
||
'type': image_type,
|
||
'omit_background': False,
|
||
'scale': 'css'
|
||
}
|
||
if image_type == 'jpeg':
|
||
screenshot_args['quality'] = quality
|
||
|
||
screenshot_bytes = await page.screenshot(**screenshot_args) # type: ignore
|
||
|
||
finally:
|
||
# 归还页面到池中，而不是直接关闭
|
||
await browser_manager.release_page(page)
|
||
|
||
# 3. 保存文件
|
||
output_path = os.path.join(self.temp_dir, output_name)
|
||
with open(output_path, "wb") as f:
|
||
f.write(screenshot_bytes)
|
||
|
||
logger.info(f"图片已生成: {output_path} ({len(screenshot_bytes)/1024:.2f} KB)")
|
||
return os.path.abspath(output_path)
|
||
|
||
except Exception as e:
|
||
logger.exception(f"渲染模板 {template_name} 失败: {e}")
|
||
return None
|
||
|
||
async def render_template_to_base64(self, template_name: str, data: Dict[str, Any], output_name: str = "output.png", quality: int = 80, image_type: str = "png") -> Optional[str]:
|
||
"""
|
||
渲染模板并返回 Base64 编码的图片字符串
|
||
"""
|
||
file_path = await self.render_template(template_name, data, output_name, quality, image_type)
|
||
if not file_path:
|
||
return None
|
||
|
||
try:
|
||
with open(file_path, "rb") as f:
|
||
content = f.read()
|
||
|
||
mime_type = "image/jpeg" if image_type == "jpeg" else "image/png"
|
||
return f"data:{mime_type};base64," + base64.b64encode(content).decode("utf-8")
|
||
except Exception as e:
|
||
logger.error(f"读取图片文件失败: {e}")
|
||
return None
|
||
|
||
# 全局图片管理器实例
|
||
image_manager = ImageManager()
|