akari-bot/modules/chemical_code/__init__.py

import asyncio
import random
import traceback

from bs4 import BeautifulSoup
from sqlalchemy import create_engine, Column, String, Text, Integer
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from tenacity import retry, stop_after_attempt

from datetime import datetime

from core.component import on_command
from core.elements import MessageSession, Image, Plain
from core.utils import get_url, download_to_cache, random_cache_path
from core.logger import Logger

from PIL import Image as PILImage


Base = declarative_base()

DB_LINK = 'sqlite:///modules/chemical_code/answer.db'  # 题型数据库


class Answer(Base):  # 数据表，为了使用 sqlalchemy 进行数据库操作所以预设此类
    __tablename__ = "Answer"
    id = Column(Integer, primary_key=True)
    cas = Column(String(512))
    answer = Column(Text)


class MSGDBSession:  # 数据库会话类
    def __init__(self):
        self.engine = engine = create_engine(DB_LINK)
        Base.metadata.create_all(bind=engine, checkfirst=True)
        self.Session = sessionmaker()
        self.Session.configure(bind=self.engine)

    @property
    def session(self):
        return self.Session()


session = MSGDBSession().session  # 实例化数据库会话并获取数据库会话


def auto_rollback_error(func):  # 函数装饰器，用于捕获异常并回滚数据库
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except Exception as e:
            session.rollback()
            raise e

    return wrapper


@retry(stop=stop_after_attempt(3))  # 函数装饰器，用于重试 3 次
@auto_rollback_error  # 函数装饰器，用于捕获异常并回滚数据库
def randcc():  # 随机从数据库（数据来源：chemicalbook）中获取一个题目
    num = random.randint(1, 20000)  # 在 1 到 20000 之间随机一个数作为抽取 ID
    query = session.query(Answer).filter_by(id=num).first()  # 根据 ID 查询题目
    return query.cas, query.answer  # 返回 chemicalbook 中的 CAS 号和化学式


csr_link = 'https://www.chemspider.com'  # ChemSpider 的链接


@retry(stop=stop_after_attempt(3), reraise=True)
async def search_csr(id=None):  # 根据 ChemSpider 的 ID 查询 ChemSpider 的链接，留空（将会使用缺省值 None）则随机查询
    if id is not None:  # 如果传入了 ID，则使用 ID 查询
        answer = id
    else:
        cas, answer = randcc()  # 否则随机查询一个题目
    get = await get_url(csr_link + '/Search.aspx?q=' + answer, 200, fmt='text')  # 在 ChemSpider 上搜索此化学式或 ID
    # Logger.info(get)
    soup = BeautifulSoup(get, 'html.parser')  # 解析 HTML
    rlist = []  # 创建一个空列表用于存放搜索结果
    try:  # 尝试获取搜索结果
        results = soup.find_all('tbody')[0].find_all('tr')  # 获取搜索结果中的所有表格行
        for x in results:  # 遍历所有表格行
            sub = x.find_all('td')[0:4]  # 获取表格行中的前四个单元格
            name = sub[2].text  # 单元格中的化学式名称
            image = sub[1].find_all('img')[0].get('src')  # 单元格中的图片链接
            rlist.append({'name': name, 'image': csr_link + image + '&w=500&h=500'})  # 将化学式名称和图片链接加入列表
    except IndexError:  # 尝试失败，进行第二次尝试
        try:
            name = soup.find('span',
                             id='ctl00_ctl00_ContentSection_ContentPlaceHolder1_RecordViewDetails_rptDetailsView_ctl00_prop_MF').text  # 获取化学式名称
            image = soup.find('img',
                              id='ctl00_ctl00_ContentSection_ContentPlaceHolder1_RecordViewDetails_rptDetailsView_ctl00_ThumbnailControl1_viewMolecule')\
                .get('src')  # 获取图片链接
            rlist.append({'name': name, 'image': csr_link + image})
        except Exception as e:  # 尝试失败，抛出错误
            raise e


    return rlist


cc = on_command('chemical_code', alias=['cc', 'chemicalcode'], desc='化学式验证码测试', developers=['OasisAkari'])
play_state = {}  # 创建一个空字典用于存放游戏状态


@cc.handle()  # 直接使用 cc 命令将触发此装饰器
async def chemical_code_by_random(msg: MessageSession):
    await c(msg)  # 将消息会话传入 c 函数


@cc.handle('stop {停止当前的游戏。}')
async def s(msg: MessageSession):
    state = play_state.get(msg.target.targetId, False)  # 尝试获取 play_state 中是否有此对象的游戏状态
    if state:  # 若有
        if state['active']:  # 检查是否为活跃状态
            play_state[msg.target.targetId]['active'] = False  # 标记为非活跃状态
            await msg.sendMessage(f'已停止，正确答案是 {state["answer"]}', quote=False)  # 发送存储于 play_state 中的答案
        else:
            await msg.sendMessage('当前无活跃状态的游戏。')
    else:
        await msg.sendMessage('当前无活跃状态的游戏。')


@cc.handle('<chemspider id> {根据 ChemSpider ID 出题}')
async def chemical_code_by_id(msg: MessageSession):
    id = msg.parsed_msg['<chemspider id>']  # 从已解析的消息中获取 ChemSpider ID
    if id.isdigit():  # 如果 ID 为纯数字
        await c(msg, id)  # 将消息会话和 ID 一并传入 c 函数
    else:
        await msg.finish('请输入纯数字ID！')


async def c(msg: MessageSession, id=None):  # 要求传入消息会话和 ChemSpider ID，ID 留空将会使用缺省值 None
    if msg.target.targetId in play_state and play_state[msg.target.targetId]['active']:  # 检查对象（群组或私聊）是否在 play_state 中有记录及是否为活跃状态
        await msg.finish('当前有一局游戏正在进行中。')
    play_state.update({msg.target.targetId: {'active': True}})  # 若无，则创建一个新的记录并标记为活跃状态
    try:
        csr = await search_csr(id)  # 尝试获取 ChemSpider ID 对应的化学式列表
    except Exception as e:  # 意外情况
        traceback.print_exc()  # 打印错误信息
        play_state[msg.target.targetId]['active'] = False  # 将对象标记为非活跃状态
        return await msg.finish('发生错误：拉取题目失败，请重新发起游戏。')
    # print(csr)
    choice = random.choice(csr)  # 从列表中随机选择一个结果
    play_state[msg.target.targetId]['answer'] = choice['name']  # 将正确答案标记于 play_state 中存储的对象中
    Logger.info(f'Answer: {choice["name"]}')  # 在日志中输出正确答案

    # 下面的代码已被注释，因为不再使用 ChemicalBook 图片数据源
    """get_image = await download_to_cache(f'https://www.chemicalbook.com/CAS/GIF/{get_rand[0]}.gif')
    Logger.info(get_rand[1])
    play_state[msg.target.targetId]['answer'] = get_rand[1]

    with PILImage.open(get_image) as im:
        if im.size[0] < 10:
            del play_state[msg.target.targetId]
            return await _(msg)
        im.seek(0)
        image = im.convert("RGBA")
        datas = image.getdata()
        newData = []
        for item in datas:
            if item[3] == 0:  # if transparent
                newData.append((230, 230, 230))  # set transparent color in jpg
            else:
                newData.append(tuple(item[:3]))
        image = PILImage.new("RGBA", im.size)
        image.getdata()
        image.putdata(newData)
        newpath = random_cache_path() + '.png'
        image.save(newpath)"""

    download = await download_to_cache(choice['image'])  # 从结果中获取链接并下载图片

    with PILImage.open(download) as im:  # 打开下载的图片
        datas = im.getdata()  # 获取图片数组
        newData = []
        for item in datas:  # 对每个像素点进行处理
            if item[3] == 0:  # 如果为透明
                newData.append((230, 230, 230))  # 设置白底
            else:
                newData.append(tuple(item[:3]))  # 否则保留原图像素点
        image = PILImage.new("RGBA", im.size)  # 创建新图片
        image.getdata()  # 获取新图片数组
        image.putdata(newData)  # 将处理后的数组覆盖新图片
        newpath = random_cache_path() + '.png'  # 创建新文件名
        image.save(newpath)  # 保存新图片

    await msg.sendMessage([Image(newpath),
                           Plain('请于2分钟内发送正确答案。（请使用字母表顺序，如：CHBrClF）')])
    time_start = datetime.now().timestamp()  # 记录开始时间

    async def ans(msg: MessageSession, answer):  # 定义回答函数的功能
        wait = await msg.waitAnyone()  # 等待对象内的任意人回答
        if play_state[msg.target.targetId]['active']:  # 检查对象是否为活跃状态
            if wait.asDisplay() != answer:  # 如果回答不正确
                Logger.info(f'{wait.asDisplay()} != {answer}')  # 输出日志
                return await ans(wait, answer)  # 进行下一轮检查
            else:
                await wait.sendMessage('回答正确。')
                play_state[msg.target.targetId]['active'] = False  # 将对象标记为非活跃状态

    async def timer(start):  # 计时器函数
        if play_state[msg.target.targetId]['active']:  # 检查对象是否为活跃状态
            if datetime.now().timestamp() - start > 120:  # 如果超过2分钟
                await msg.sendMessage(f'已超时，正确答案是 {play_state[msg.target.targetId]["answer"]}', quote=False)
                play_state[msg.target.targetId]['active'] = False
            else:  # 如果未超时
                await asyncio.sleep(1)  # 等待1秒
                await timer(start)  # 重新调用计时器函数

    await asyncio.gather(ans(msg, choice['name']), timer(time_start))  # 同时启动回答函数和计时器函数