# -*- coding: utf-8 -*- import traceback from typing import Dict from aiohttp.client_exceptions import ClientError from nonebot.plugin import on_command, on_message from nonebot.adapters.cqhttp import Bot, MessageEvent, GroupMessageEvent from nonebot.typing import T_State from services.log import logger from util.utils import get_message_text, get_message_imgs from configs.config import MAX_FIND_IMG_COUNT from .ex import get_des as get_des_ex from .iqdb import get_des as get_des_iqdb from .saucenao import get_des as get_des_sau from .ascii2d import get_des as get_des_asc from .trace import get_des as get_des_trace from .yandex import get_des as get_des_yandex async def get_des(url: str, mode: str, user_id: int): """ :param url: 图片链接 :param mode: 图源 :return: """ if mode == "iqdb": async for msg in get_des_iqdb(url): yield msg elif mode == "ex": async for msg in get_des_ex(url): yield msg elif mode == "trace": async for msg in get_des_trace(url): yield msg elif mode == "yandex": async for msg in get_des_yandex(url): yield msg elif mode.startswith("asc"): async for msg in get_des_asc(url, user_id): yield msg else: async for msg in get_des_sau(url, user_id): yield msg setu = on_command("识图", aliases={"search"}, block=True, priority=5) @setu.handle() async def handle_first_receive(bot: Bot, event: MessageEvent, state: T_State): msg = get_message_text(event.json()) imgs = get_message_imgs(event.json()) if msg in ['帮助']: await setu.finish('示例:\n\t识图 (图片)\n\t识图asc (图片)') if imgs: state["setu"] = imgs[0] if msg: state["mod"] = msg # ex/nao/trace/iqdb/ascii2d # @setu.got("mod", prompt="从哪里查找呢? ex/nao/trace/iqdb/ascii2d") # async def get_func(bot: Bot, event: MessageEvent, state: dict): # pass @setu.args_parser async def get_setu(bot: Bot, event: MessageEvent, state: T_State): imgs = get_message_imgs(event.json()) msg = get_message_text(event.json()) if not imgs: await setu.reject() if msg: state['mod'] = msg state["setu"] = imgs[0] @setu.got("setu", prompt="图呢?") async def get_setu(bot: Bot, event: MessageEvent, state: T_State): """ 发现没有的时候要发问 :return: """ url: str = state["setu"] mod: str = state["mod"] if state.get("mod") else "nao" # 模式 try: await bot.send(event=event, message="正在处理图片") idx = 1 async for msg in get_des(url, mod, event.user_id): await bot.send(event=event, message=msg) if idx == MAX_FIND_IMG_COUNT: break idx += 1 logger.info(f"(USER {event.user_id}, GROUP " f"{event.group_id if event.message_type != 'private' else 'private'}) 识图:{url}") # image_data: List[Tuple] = await get_pic_from_url(url) # await setu.finish("hso") except IndexError: # await bot.send(event, traceback.format_exc()) await setu.finish("参数错误") except ClientError: await setu.finish("连接失败") pic_map: Dict[str, str] = {} # 保存这个群的其阿金一张色图 {"123456":http://xxx"} async def check_pic(bot: Bot, event: MessageEvent, state: T_State) -> bool: if isinstance(event, MessageEvent): for msg in event.message: if msg.type == "image": url: str = msg.data["url"] state["url"] = url return True return False notice_pic = on_message(check_pic, block=False, priority=1) @notice_pic.handle() async def handle_pic(bot: Bot, event: GroupMessageEvent, state: T_State): try: group_id: str = str(event.group_id) pic_map.update({group_id: state["url"]}) except AttributeError: pass previous = on_command("上一张图是什么", aliases={"上一张", "这是什么"}) @previous.handle() async def handle_previous(bot: Bot, event: GroupMessageEvent, state: T_State): await bot.send(event=event, message="processing...") try: url: str = pic_map[str(event.group_id)] idx = 1 async for msg in get_des(url, "nao", event.user_id): await bot.send(event=event, message=msg) if idx == MAX_FIND_IMG_COUNT: break idx += 1 except IndexError: await previous.finish("参数错误") except ClientError: await previous.finish("连接错误") except KeyError: await previous.finish("没有图啊QAQ")