import os
import tempfile
import subprocess
import base64
from typing import Any, AsyncGenerator
from astrbot.api.event import filter, AstrMessageEvent, MessageEventResult
from astrbot.api.star import Context, Star
from astrbot.api import logger
from astrbot.core.message.message_event_result import MessageEventResult
class SuanPlugin(Star):
def __init__(self, context: Context):
super().__init__(context)
async def initialize(self):
"""可选择实现异步的插件初始化方法,当实例化该插件类之后会自动调用该方法。"""
@filter.llm_tool(name='render_course_schedule_png')
async def render_course_schedule_png(self, event: AstrMessageEvent, schedule: list[tuple], title: str) -> \
AsyncGenerator[MessageEventResult, Any]:
"""将完整的每周课表渲染为 png 格式的图片,或者说是得到一张课程表图片。
在调用之前可能需要先对 schedule 进行格式转换,本工具接收的数据类型是 list[tuple]。此外,应当尽可能将代表周数的元素从列表压缩为字符串以节约工具参数长度。
Args:
schedule (list[tuple]): 包含数个元组的列表。每个元组的内部是按照顺序的如下元素:
0. str: 课程名称
1. str: 授课教师(可以为 None)
2. str: 教室(可以为 None)
3. str | list: 形如 `1-17` 表示从第 1 周到第 17 周;`1-3,5,7-17` 也是受支持的;list[int] 也是受支持的。
4. int: 星期几,int (1-7)
5. list: 当日第几节课,一般是相邻的两节,list[int]
title (str): 图片名称,会被渲染在图片的顶部,格式上作为课程表标题。
"""
# 转换为字典格式
schedule_dict = convert_tuple_schedule_to_dict(schedule)
try:
# 生成 HTML 内容
html_content = self._generate_schedule_html(schedule_dict, title)
# 创建临时文件
with tempfile.NamedTemporaryFile(mode='w', suffix='.html', delete=False, encoding='utf-8') as f:
f.write(html_content)
html_path = f.name
# 生成 PNG 输出路径
png_path = html_path.replace('.html', '.png')
# 使用 wkhtmltoimage 直接生成图片(比 wkhtmltopdf 再转换更方便)
# 如果系统只有 wkhtmltopdf,可以先生成 PDF 再转换
try:
# 尝试使用 wkhtmltoimage(如果安装了)
cmd = [
'wkhtmltoimage',
'--width', '1200',
'--quality', '90',
'--enable-local-file-access',
html_path,
png_path
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode != 0 or not os.path.exists(png_path):
# 回退到 wkhtmltopdf + convert 方案
png_path = await self._html_to_png_via_pdf(html_path)
except FileNotFoundError:
# wkhtmltoimage 不存在,使用 pdf 方案
png_path = await self._html_to_png_via_pdf(html_path)
# 读取图片并返回
if os.path.exists(png_path):
with open(png_path, 'rb') as f:
image_data = f.read()
# 清理临时文件
self._cleanup_temp_files(html_path, png_path)
# 返回图片消息
from astrbot.api.message_components import Image
yield event.chain_result([Image.fromBytes(image_data)])
else:
yield event.plain_result("生成课表图片失败,请检查 wkhtmltopdf 是否已安装")
except Exception as e:
logger.error(f"渲染课表失败: {e}")
yield event.plain_result(f"生成课表图片时出错: {str(e)}")
async def _html_to_png_via_pdf(self, html_path: str) -> str:
"""通过 PDF 中转生成 PNG"""
pdf_path = html_path.replace('.html', '.pdf')
png_path = html_path.replace('.html', '.png')
# HTML -> PDF
cmd = [
'wkhtmltopdf',
'--page-width', '297mm', # A4 横向
'--page-height', '210mm',
'--encoding', 'utf-8',
'--enable-local-file-access',
'--no-stop-slow-scripts',
'--javascript-delay', '1000',
html_path,
pdf_path
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode != 0:
raise RuntimeError(f"wkhtmltopdf 失败: {result.stderr}")
# PDF -> PNG (使用 ImageMagick 的 convert)
cmd = [
'convert',
'-density', '150',
pdf_path,
'-quality', '90',
png_path
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode != 0:
# 尝试使用 pdftoppm
cmd = [
'pdftoppm',
'-png',
'-r', '150',
'-singlefile',
pdf_path,
png_path.replace('.png', '')
]
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
if result.returncode != 0:
raise RuntimeError(f"PDF 转 PNG 失败: {result.stderr}")
# pdftoppm 生成的文件名可能不同
generated_png = png_path.replace('.png', '-1.png')
if os.path.exists(generated_png):
os.rename(generated_png, png_path)
# 清理 PDF 临时文件
if os.path.exists(pdf_path):
os.unlink(pdf_path)
return png_path
def _cleanup_temp_files(self, html_path: str, png_path: str):
"""清理临时文件"""
for path in [html_path, png_path]:
try:
if os.path.exists(path):
os.unlink(path)
except Exception as e:
logger.warning(f"清理临时文件失败 {path}: {e}")
def _generate_schedule_html(self, schedule: list[dict], title: str) -> str:
"""生成课程表 HTML"""
# 星期映射
weekdays = ['周一', '周二', '周三', '周四', '周五', '周六', '周日']
# 初始化课表网格: 12节课 x 7天
# grid[节次][星期] = 课程列表
grid = [[[] for _ in range(7)] for _ in range(12)]
# 填充课程数据
for course in schedule:
name = course.get('name', '未知课程')
teacher = course.get('teacher') or ''
classroom = course.get('classroom') or ''
weeks = course.get('weeks', [])
day = course.get('day', 1) # 1-7
classes = course.get('classes', [])
# 调整 day 为 0-based 索引
day_idx = day - 1
if day_idx < 0 or day_idx > 6:
continue
# 构建课程显示文本
course_text = f"{name}"
if teacher:
course_text += f"
{teacher}"
if classroom:
course_text += f"
@{classroom}"
if weeks:
course_text += f"
第{self._format_weeks(weeks)}周"
# 放入对应格子
for class_idx in classes:
if 1 <= class_idx <= 12:
grid[class_idx - 1][day_idx].append(course_text)
# 构建 HTML
html = f'''
| 节次 | {''.join(f'{w} | ' for w in weekdays)}
|---|