Files
NeoBot/core/managers/image_manager.py
K2cr2O1 5f943c1792 feat(github_parser): 添加GitHub仓库信息查询功能
- 新增github_parser插件,支持通过命令或自动解析链接查询GitHub仓库信息
- 添加github_repo.html模板用于渲染仓库信息图片
- 优化图片管理器支持高质量截图和CSS缩放
- 重构消息事件类权限常量定义方式
- 更新帮助页面样式为三列布局并优化响应式设计
2026-01-20 18:33:46 +08:00

130 lines
5.0 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
图片生成管理器模块
负责管理图片生成相关的逻辑,支持多种渲染引擎(目前支持 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
class ImageManager:
"""
图片生成管理器(单例)
"""
_instance = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
# 模板目录
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()