summaryrefslogtreecommitdiff
path: root/ATRI/plugins/saucenao
diff options
context:
space:
mode:
Diffstat (limited to 'ATRI/plugins/saucenao')
-rw-r--r--ATRI/plugins/saucenao/__init__.py7
-rw-r--r--ATRI/plugins/saucenao/data_source.py25
2 files changed, 21 insertions, 11 deletions
diff --git a/ATRI/plugins/saucenao/__init__.py b/ATRI/plugins/saucenao/__init__.py
index 468b001..b55f18c 100644
--- a/ATRI/plugins/saucenao/__init__.py
+++ b/ATRI/plugins/saucenao/__init__.py
@@ -16,6 +16,7 @@ _search_flmt_notice = choice(["慢...慢一..点❤", "冷静1下", "歇会歇�
saucenao = SaouceNao().on_command("以图搜图", "透过一张图搜索可能的来源")
+
@saucenao.args_parser # type: ignore
async def _get_img(bot: Bot, event: MessageEvent, state: T_State):
msg = str(event.message).strip()
@@ -27,16 +28,18 @@ async def _get_img(bot: Bot, event: MessageEvent, state: T_State):
else:
state["img"] = msg
+
@saucenao.handle()
async def _ready_search(bot: Bot, event: MessageEvent, state: T_State):
user_id = event.get_user_id()
if not _search_flmt.check(user_id):
await saucenao.finish(_search_flmt_notice)
-
+
msg = str(event.message).strip()
if msg:
state["img"] = msg
+
@saucenao.got("img", "图呢?")
async def _deal_search(bot: Bot, event: MessageEvent, state: T_State):
user_id = event.get_user_id()
@@ -44,7 +47,7 @@ async def _deal_search(bot: Bot, event: MessageEvent, state: T_State):
img = findall(r"url=(.*?)]", msg)
if not img:
await saucenao.reject("请发送图片而不是其他东西!!")
-
+
a = SaouceNao(SauceNAO.key)
result = f"> {MessageSegment.at(user_id)}" + await a.search(img[0])
_search_flmt.start_cd(user_id)
diff --git a/ATRI/plugins/saucenao/data_source.py b/ATRI/plugins/saucenao/data_source.py
index d948f91..6e72c98 100644
--- a/ATRI/plugins/saucenao/data_source.py
+++ b/ATRI/plugins/saucenao/data_source.py
@@ -16,8 +16,15 @@ __doc__ = """
class SaouceNao(Service):
-
- def __init__(self, api_key: str = None, output_type=2, testmode=1, dbmaski=32768, db=5, numres=5):
+ def __init__(
+ self,
+ api_key: str = None,
+ output_type=2,
+ testmode=1,
+ dbmaski=32768,
+ db=5,
+ numres=5,
+ ):
Service.__init__(self, "以图搜图", __doc__, rule=is_in_service("以图搜图"))
params = dict()
@@ -28,31 +35,31 @@ class SaouceNao(Service):
params["db"] = db
params["numres"] = numres
self.params = params
-
+
async def _request(self, url: str):
self.params["url"] = url
-
+
try:
res = await request.post(URL, params=self.params)
except RequestError:
raise RequestError("Request failed!")
data = await res.json()
return data
-
+
async def search(self, url: str) -> str:
data = await self._request(url)
res = data["results"]
-
+
result = list()
for i in range(len(res)):
data = res[i]
-
+
_result = dict()
_result["similarity"] = data["header"]["similarity"]
_result["index_name"] = data["header"]["index_name"]
_result["url"] = choice(data["data"].get("ext_urls", ["None"]))
result.append(_result)
-
+
msg0 = str()
for i in result:
msg0 += (
@@ -61,7 +68,7 @@ class SaouceNao(Service):
f"Name: {i['index_name']}\n"
f"URL: {i['url'].replace('https://', '')}"
)
-
+
if len(res) <= 3:
return msg0
else: