diff options
Diffstat (limited to 'ATRI/plugins/anime_search.py')
| -rw-r--r-- | ATRI/plugins/anime_search.py | 167 |
1 files changed, 91 insertions, 76 deletions
diff --git a/ATRI/plugins/anime_search.py b/ATRI/plugins/anime_search.py index 6bb11e8..de64501 100644 --- a/ATRI/plugins/anime_search.py +++ b/ATRI/plugins/anime_search.py @@ -1,65 +1,48 @@ import re -import json from aiohttp import FormData +from random import choice -from nonebot.adapters.cqhttp import Bot, MessageEvent -from nonebot.adapters.cqhttp.message import Message from nonebot.typing import T_State +from nonebot.adapters.cqhttp import Bot, MessageEvent +from nonebot.adapters.cqhttp.message import Message, MessageSegment -from ATRI.service import Service as sv +from ATRI.service import Service from ATRI.rule import is_in_service +from ATRI.utils import request, UbuntuPaste, Translate +from ATRI.utils.limit import FreqLimiter from ATRI.exceptions import RequestError -from ATRI.utils.request import get_bytes -from ATRI.utils.translate import to_simple_string -from ATRI.utils.ub_paste import paste URL = "https://trace.moe/api/search?url=" +_anime_flmt = FreqLimiter(10) +_anime_flmt_notice = choice(["慢...慢一..点❤", "冷静1下", "歇会歇会~~"]) __doc__ = """ -以图搜番 -权限组：所有人 -用法： - 以图搜番 (pic) +通过一张图片搜索你需要的番！据说里*也可以 """ -anime_search = sv.on_command(cmd="以图搜番", docs=__doc__, rule=is_in_service("以图搜番")) - - -@anime_search.args_parser # type: ignore -async def _load_anime(bot: Bot, event: MessageEvent, state: T_State) -> None: - msg = str(event.message) - quit_list = ["算了", "罢了", "不搜了", "取消"] - if msg in quit_list: - await anime_search.finish("好吧...") - if not msg: - await anime_search.reject("图呢？") - else: - state["pic_anime"] = msg - - -@anime_search.handle() -async def _anime_search(bot: Bot, event: MessageEvent, state: T_State) -> None: - msg = str(event.message).strip() - if msg: - state["pic_anime"] = msg - - -@anime_search.got("pic_anime", prompt="图呢？") -async def _deal_search(bot: Bot, event: MessageEvent, state: T_State) -> None: - msg = state["pic_anime"] - img = re.findall(r"url=(.*?)]", msg) - if not img: - await anime_search.reject("请发送图片而不是其它东西！！") - - try: - req = await get_bytes(URL + img[0]) - except RequestError: - raise RequestError("Request failed!") - data = json.loads(req)["docs"] - try: +class Anime(Service): + + def __init__(self): + Service.__init__(self, "以图搜番", __doc__, rule=is_in_service("以图搜番")) + + @staticmethod + async def _request(url: str) -> dict: + aim = URL + url + try: + res = await request.get(aim) + except RequestError: + raise RequestError("Request failed!") + result = await res.json() + return result + + @classmethod + async def search(cls, url: str) -> str: + data = await cls._request(url) + data = data["docs"] + d = dict() for i in range(len(data)): if data[i]["title_chinese"] in d.keys(): @@ -73,38 +56,70 @@ async def _deal_search(bot: Bot, event: MessageEvent, state: T_State) -> None: else: n = data[i]["episode"] - d[to_simple_string(data[i]["title_chinese"])] = [ + d[Translate(data[i]["title_chinese"]).to_simple()] = [ data[i]["similarity"], f"第{n}集", f"{int(m)}分{int(s)}秒处", ] - except Exception as err: - raise Exception(f"Invalid data.\n{err}") - - result = sorted(d.items(), key=lambda x: x[1], reverse=True) - - t = 0 - - msg0 = f"> {event.sender.nickname}" - for i in result: - t += 1 - s = "%.2f%%" % (i[1][0] * 100) - msg0 = msg0 + ( - "\n——————————\n" - f"({t}) Similarity: {s}\n" - f"Name: {i[0]}\n" - f"Time: {i[1][1]} {i[1][2]}" - ) - - if len(result) == 2: - await anime_search.finish(Message(msg0)) + + result = sorted(d.items(), key=lambda x: x[1], reverse=True) + t = 0 + msg0 = str() + for i in result: + t += 1 + s = "%.2f%%" % (i[1][0] * 100) + msg0 = msg0 + ( + "\n——————————\n" + f"({t}) Similarity: {s}\n" + f"Name: {i[0]}\n" + f"Time: {i[1][1]} {i[1][2]}" + ) + + if len(result) == 2: + return msg0 + else: + data = FormData() + data.add_field("poster", "ATRI running log") + data.add_field("syntax", "text") + data.add_field("expiration", "day") + data.add_field("content", msg0) + + repo = f"详细请移步此处~\n{await UbuntuPaste(data).paste()}" + return repo + + +anime_search = Anime().on_command("以图搜番", "发送一张图以搜索可能的番剧") + +@anime_search.args_parser # type: ignore +async def _get_anime(bot: Bot, event: MessageEvent, state: T_State): + msg = str(event.message).strip() + quit_list = ["算了", "罢了", "不搜了", "取消"] + if msg in quit_list: + await anime_search.finish("好吧...") + if not msg: + await anime_search.reject("图呢？") else: - data = FormData() - data.add_field("poster", "ATRI running log") - data.add_field("syntax", "text") - data.add_field("expiration", "day") - data.add_field("content", msg0) - - repo = f"> {event.sender.nickname}\n" - repo = repo + f"详细请移步此处~\n{await paste(data)}" - await anime_search.finish(repo) + state["anime"] = msg + +@anime_search.handle() +async def _ready_sear(bot: Bot, event: MessageEvent, state: T_State): + user_id = event.get_user_id() + if not _anime_flmt.check(user_id): + await anime_search.finish(_anime_flmt_notice) + + msg = str(event.message).strip() + if msg: + state["anime"] = msg + +@anime_search.got("anime", "图呢？") +async def _deal_sear(bot: Bot, event: MessageEvent, state: T_State): + user_id = event.get_user_id() + msg = state["anime"] + img = re.findall(r"url=(.*?)]", msg) + if not img: + await anime_search.reject("请发送图片而不是其它东西！！") + + a = await Anime().search(img[0]) + result = f"> {MessageSegment.at(user_id)}\n" + a + _anime_flmt.start_cd(user_id) + await anime_search.finish(Message(result)) |