From fdebf4d60f5d2f38669182fdafa0ca34a37d0365 Mon Sep 17 00:00:00 2001 From: Eigeen Date: Mon, 27 Mar 2023 21:19:36 +0800 Subject: [PATCH] =?UTF-8?q?=E9=A2=84=E8=A7=88=E5=9B=BE=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E9=80=9A=E8=BF=87=E4=BB=A3=E7=90=86=E4=B8=8B=E8=BD=BD=EF=BC=8C?= =?UTF-8?q?=E5=8F=AF=E9=80=89=E4=B8=8D=E6=98=BE=E7=A4=BA=E9=A2=84=E8=A7=88?= =?UTF-8?q?=E5=9B=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 8 +- nonebot_plugin_hikarisearch/config.py | 15 +++ nonebot_plugin_hikarisearch/data_source.py | 115 ++++++++++++++------- 3 files changed, 98 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index a7ddf18..efe3168 100644 --- a/README.md +++ b/README.md @@ -43,7 +43,9 @@ pip install nonebot_plugin_hikarisearch 可在 `.env.xxx` 文件中添加如下配置： ``` -hikarisearch_api=xxx # HikariSearch 站点，默认为 "https://hikari.obfs.dev" -hikarisearch_max_results=xxx # 最多返回的结果数量，默认为 3 -hikarisearch_withdraw=xxx # 自动撤回时间，默认为 0 (不撤回)，单位为秒 +hikarisearch_api=xxx # HikariSearch 站点，默认为 "https://hikari.obfs.dev" +hikarisearch_max_results=3 # 最多返回的结果数量，默认为 3 +hikarisearch_withdraw=0 # 自动撤回时间，默认为 0 (不撤回)，单位为秒 +hikarisearch_proxy=xxx # 代理，适用于搜索结果预览图下载 +hikarisearch_thumb=true # 是否发送搜索结果预览图 ``` diff --git a/nonebot_plugin_hikarisearch/config.py b/nonebot_plugin_hikarisearch/config.py index 07f8c72..da9f1d3 100644 --- a/nonebot_plugin_hikarisearch/config.py +++ b/nonebot_plugin_hikarisearch/config.py @@ -6,6 +6,21 @@ class Config(BaseModel, extra=Extra.ignore): hikarisearch_api: str = "https://hikari.obfs.dev" hikarisearch_max_results: int = 3 hikarisearch_withdraw: int = 0 + _hikarisearch_proxy: str = '' + hikarisearch_thumb: bool = True + + # hikarisearch_proxy只能为None或非空字符串 + @property + def hikarisearch_proxy(self): + if self._hikarisearch_proxy == '': + return None + return self._hikarisearch_proxy + + @hikarisearch_proxy.setter + def hikarisearch_proxy(self, value): + if value is None: + value = '' + self._hikarisearch_proxy = value hikari_config = Config.parse_obj(get_driver().config.dict()) diff --git a/nonebot_plugin_hikarisearch/data_source.py b/nonebot_plugin_hikarisearch/data_source.py index 9d11ee0..c38559b 100644 --- a/nonebot_plugin_hikarisearch/data_source.py +++ b/nonebot_plugin_hikarisearch/data_source.py @@ -1,7 +1,8 @@ -import httpx -from dataclasses import dataclass from typing import List, Tuple, Protocol +from dataclasses import dataclass +import httpx from nonebot.adapters.onebot.v11 import Message, MessageSegment +from nonebot.log import logger from .config import hikari_config @@ -11,17 +12,24 @@ async def search_saucenao(image: bytes) -> List[Message]: data = {"hide": "true"} result = await post("/api/SauceNAO", data, image) - return [ - MessageSegment.image(res["image"]) - + "{}\n图片相似度: {:.2f}%\n图片来源:\n{}".format( - res["title"], - res["similarity"], - "\n".join( - ["\n".join(dict(content).values()) for content in res["content"]] - ), - ) - for res in result - ] + + msg = [] + for res in result: + _msg = '' + if hikari_config.hikarisearch_thumb: + _msg += await thumb_msg(res["image"]) + + _msg += "{}\n图片相似度: {:.2f}%\n图片来源:\n{}".format( + res["title"], + res["similarity"], + "\n".join( + ["\n".join(dict(content).values()) + for content in res["content"]] + ), + ) + msg.append(_msg) + + return msg async def search_iqdb(image: bytes) -> List[Message]: @@ -30,67 +38,100 @@ async def search_iqdb(image: bytes) -> List[Message]: "discolor": "false", } result = await post("/api/IqDB", data, image) - return [ - MessageSegment.image(res["image"]) - + "图片相似度: {:.0f}%\n图片来源:\n{}".format(res["similarity"], res["url"]) - for res in result - ] + + msg = [] + for res in result: + _msg = '' + if hikari_config.hikarisearch_thumb: + _msg += await thumb_msg(res["image"]) + + _msg += "图片相似度: {:.0f}%\n图片来源:\n{}".format(res["similarity"], res["url"]) + msg.append(_msg) + + return msg async def search_ascii2d(image: bytes) -> List[Message]: data = {"type": "color"} result = await post("/api/ascii2d", data, image) - return [ - MessageSegment.image(res["image"]) - + "图片来源: {}\n{}\n图片作者: {}\n{}".format( + + msg = [] + for res in result: + _msg = '' + if hikari_config.hikarisearch_thumb: + _msg += await thumb_msg(res["image"]) + + _msg += "图片来源: {}\n{}\n图片作者: {}\n{}".format( res["source"]["text"], res["source"]["link"], res["author"]["text"], res["author"]["link"], ) - for res in result - ] + msg.append(_msg) + + return msg async def search_ehentai(image: bytes) -> List[Message]: data = {"site": "eh", "cover": "false", "deleted": "false", "similar": "true"} result = await post("/api/E-Hentai", data, image) - return [ - MessageSegment.image(res["image"]) - + "[{}] {}\n图片来源:\n{}".format(res["type"], res["title"], res["link"]) - for res in result - ] + + msg = [] + for res in result: + _msg = '' + if hikari_config.hikarisearch_thumb: + _msg += await thumb_msg(res["image"]) + + _msg += "[{}] {}\n图片来源:\n{}".format(res["type"], res["title"], res["link"]) + msg.append(_msg) + + return msg async def search_tracemoe(image: bytes) -> List[Message]: data = {"cutBorders": "true"} result = await post("/api/TraceMoe", data, image) - return [ - MessageSegment.image(res["preview"]) - + "图片相似度:{:.2f}\n{}图片来源:\n{}\n英文名:{}\n罗马字名:{}".format( + + msg = [] + for res in result: + _msg = '' + if hikari_config.hikarisearch_thumb: + _msg += await thumb_msg(res["image"]) + + _msg += "图片相似度:{:.2f}\n{}图片来源:\n{}\n英文名:{}\n罗马字名:{}".format( res["similarity"], res["file"], res["name"]["native"], res["name"]["english"], res["name"]["romaji"], ) - for res in result - ] + msg.append(_msg) + + return msg async def post(url, data: dict, image: bytes) -> dict: files = {"image": image} - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(proxies=hikari_config.hikarisearch_proxy) as client: resp = await client.post(API + url, data=data, files=files, timeout=20) return resp.json() -async def download_image(url: str) -> bytes: - async with httpx.AsyncClient() as client: - resp = await client.get(url, timeout=20) +async def download_image(url: str, timeout=20) -> bytes: + async with httpx.AsyncClient(proxies=hikari_config.hikarisearch_proxy) as client: + resp = await client.get(url, timeout=timeout) return resp.content +async def thumb_msg(url: str) -> MessageSegment: + try: + img = await download_image(url, timeout=5) + return MessageSegment.image(img) + except Exception as e: + logger.warning('下载搜索结果预览图失败，将使用原始图片链接。原因: {}', e) + return MessageSegment.image(url) + + class Func(Protocol): async def __call__(self, image: bytes) -> List[Message]: ...