443 lines
14 KiB
Python
443 lines
14 KiB
Python
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"<b>{name}</b>"
|
||
if teacher:
|
||
course_text += f"<br><small>{teacher}</small>"
|
||
if classroom:
|
||
course_text += f"<br><small>@{classroom}</small>"
|
||
if weeks:
|
||
course_text += f"<br><small>第{self._format_weeks(weeks)}周</small>"
|
||
|
||
# 放入对应格子
|
||
for class_idx in classes:
|
||
if 1 <= class_idx <= 12:
|
||
grid[class_idx - 1][day_idx].append(course_text)
|
||
|
||
# 构建 HTML
|
||
html = f'''<!DOCTYPE html>
|
||
<html>
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<title>{title}</title>
|
||
<style>
|
||
* {{
|
||
margin: 0;
|
||
padding: 0;
|
||
box-sizing: border-box;
|
||
}}
|
||
body {{
|
||
font-family: "Microsoft YaHei", "SimHei", sans-serif;
|
||
padding: 20px;
|
||
background: #fff;
|
||
}}
|
||
.container {{
|
||
width: 100%;
|
||
max-width: 1200px;
|
||
margin: 0 auto;
|
||
}}
|
||
h1 {{
|
||
text-align: center;
|
||
font-size: 24px;
|
||
margin-bottom: 20px;
|
||
color: #333;
|
||
}}
|
||
table {{
|
||
width: 100%;
|
||
border-collapse: collapse;
|
||
font-size: 12px;
|
||
}}
|
||
th, td {{
|
||
border: 1px solid #ccc;
|
||
padding: 8px;
|
||
text-align: center;
|
||
vertical-align: middle;
|
||
}}
|
||
th {{
|
||
background: #f5f5f5;
|
||
font-weight: bold;
|
||
color: #333;
|
||
}}
|
||
.time-col {{
|
||
background: #f9f9f9;
|
||
font-weight: bold;
|
||
width: 60px;
|
||
}}
|
||
.course-cell {{
|
||
min-height: 60px;
|
||
}}
|
||
.course-item {{
|
||
background: #e3f2fd;
|
||
border-radius: 4px;
|
||
padding: 4px;
|
||
margin: 2px 0;
|
||
line-height: 1.4;
|
||
}}
|
||
.course-item b {{
|
||
color: #1976d2;
|
||
font-size: 11px;
|
||
}}
|
||
.course-item small {{
|
||
color: #666;
|
||
font-size: 9px;
|
||
}}tr:nth-child(even) td {{
|
||
background: #fafafa;
|
||
}}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<div class="container">
|
||
<h1>{title}</h1>
|
||
<table>
|
||
<thead>
|
||
<tr>
|
||
<th>节次</th>
|
||
{''.join(f'<th>{w}</th>' for w in weekdays)}
|
||
</tr>
|
||
</thead>
|
||
<tbody>
|
||
{self._generate_table_body(grid)}
|
||
</tbody>
|
||
</table>
|
||
</div>
|
||
</body>
|
||
</html>'''
|
||
|
||
return html
|
||
|
||
def _generate_table_body(self, grid: list) -> str:
|
||
"""生成表格主体 HTML"""
|
||
rows = []
|
||
|
||
for period in range(12):
|
||
# 计算时间显示
|
||
start_time = self._get_period_time(period + 1)
|
||
|
||
row_html = f'<tr><td class="time-col">{period + 1}<br><small>{start_time}</small></td>'
|
||
|
||
for day in range(7):
|
||
courses = grid[period][day]
|
||
if courses:
|
||
cell_content = ''.join(
|
||
f'<div class="course-item">{c}</div>'
|
||
for c in courses
|
||
)
|
||
else:
|
||
cell_content = ''
|
||
|
||
row_html += f'<td class="course-cell">{cell_content}</td>'
|
||
|
||
row_html += '</tr>'
|
||
rows.append(row_html)
|
||
|
||
return '\n'.join(rows)
|
||
|
||
def _get_period_time(self, period: int) -> str:
|
||
"""获取节次对应的时间（南邮标准时间）"""
|
||
# 南邮标准作息时间，可根据需要调整
|
||
time_map = {
|
||
1: '08:00',
|
||
2: '08:50',
|
||
3: '09:50',
|
||
4: '10:40',
|
||
5: '11:30',
|
||
6: '13:45',
|
||
7: '14:35',
|
||
8: '15:35',
|
||
9: '16:25',
|
||
10: '18:30',
|
||
11: '19:20',
|
||
12: '20:10'
|
||
}
|
||
return time_map.get(period, '')
|
||
|
||
def _format_weeks(self, weeks: list[int]) -> str:
|
||
"""格式化周数显示，如 [1,2,3,5,6,7] -> "1-3,5-7" """
|
||
if not weeks:
|
||
return ''
|
||
|
||
weeks = sorted(set(weeks))
|
||
ranges = []
|
||
start = end = weeks[0]
|
||
|
||
for w in weeks[1:]:
|
||
if w == end + 1:
|
||
end = w
|
||
else:
|
||
if start == end:
|
||
ranges.append(str(start))
|
||
else:
|
||
ranges.append(f"{start}-{end}")
|
||
start = end = w
|
||
|
||
if start == end:
|
||
ranges.append(str(start))
|
||
else:
|
||
ranges.append(f"{start}-{end}")
|
||
|
||
return ','.join(ranges)
|
||
|
||
|
||
async def terminate(self):
|
||
"""可选择实现异步的插件销毁方法，当插件被卸载/停用时会调用。"""
|
||
|
||
|
||
def convert_tuple_schedule_to_dict(schedule: list[tuple]) -> list[dict]:
|
||
"""将元组格式的课表转换为字典格式。
|
||
|
||
Args:
|
||
schedule: list[tuple]，每个元组包含 (name, teacher, classroom, weeks, day, classes)
|
||
其中 weeks 可以是 list[int] 或 str（如 "1-17"）
|
||
|
||
Returns:
|
||
list[dict]: 标准格式的课程数据
|
||
"""
|
||
result = []
|
||
for course in schedule:
|
||
# 确保元组长度足够，不足的补 None
|
||
padded = list(course) + [None] * (6 - len(course))
|
||
name, teacher, classroom, weeks, day, classes = padded[:6]
|
||
|
||
# 处理 weeks：支持 list 或 str（如 "1-17", "1,3,5-10"）
|
||
weeks_list = []
|
||
if isinstance(weeks, list):
|
||
weeks_list = weeks
|
||
elif isinstance(weeks, str):
|
||
weeks_list = parse_weeks_string(weeks)
|
||
|
||
result.append({
|
||
"name": name,
|
||
"teacher": teacher,
|
||
"classroom": classroom,
|
||
"weeks": weeks_list,
|
||
"day": day if isinstance(day, int) else 1,
|
||
"classes": classes if isinstance(classes, list) else []
|
||
})
|
||
|
||
return result
|
||
|
||
|
||
def parse_weeks_string(weeks_str: str) -> list[int]:
|
||
"""将周数字符串解析为整数列表。
|
||
|
||
支持格式：
|
||
"1-17" -> [1,2,3,...,17]
|
||
"1,3,5" -> [1,3,5]
|
||
"1-5,7,9-11" -> [1,2,3,4,5,7,9,10,11]
|
||
|
||
Args:
|
||
weeks_str: 周数字符串
|
||
|
||
Returns:
|
||
list[int]: 周数列表
|
||
"""
|
||
if not weeks_str or not isinstance(weeks_str, str):
|
||
return []
|
||
|
||
result = []
|
||
parts = weeks_str.split(',')
|
||
|
||
for part in parts:
|
||
part = part.strip()
|
||
if '-' in part:
|
||
# 范围格式：1-17
|
||
try:
|
||
start, end = part.split('-', 1)
|
||
start = int(start.strip())
|
||
end = int(end.strip())
|
||
result.extend(range(start, end + 1))
|
||
except ValueError:
|
||
continue
|
||
else:
|
||
# 单个数字
|
||
try:
|
||
result.append(int(part))
|
||
except ValueError:
|
||
continue
|
||
|
||
# 去重并排序
|
||
return sorted(set(result))
|
||
|