diff options
Diffstat (limited to 'ATRI/plugins/saucenao/data_source.py')
-rw-r--r-- | ATRI/plugins/saucenao/data_source.py | 96 |
1 files changed, 47 insertions, 49 deletions
diff --git a/ATRI/plugins/saucenao/data_source.py b/ATRI/plugins/saucenao/data_source.py index 004d9de..e270db1 100644 --- a/ATRI/plugins/saucenao/data_source.py +++ b/ATRI/plugins/saucenao/data_source.py @@ -1,69 +1,67 @@ -from random import choice +import json +from typing import List -from ATRI.exceptions import RequestError from ATRI.utils import request +from ATRI.message import MessageBuilder +from .models import SauceNAORequest, SauceNAOResponse, SauceNAOResult -URL = "https://saucenao.com/search.php" +SAUCENAO_URL: str = "https://saucenao.com/search.php" -class SauceNao: +class SauceNAO: def __init__( self, - api_key: str = str(), - output_type=2, - testmode=1, - dbmaski=32768, - db=5, - numres=5, - ): - params = dict() - params["api_key"] = api_key - params["output_type"] = output_type - params["testmode"] = testmode - params["dbmaski"] = dbmaski - params["db"] = db - params["numres"] = numres - self.params = params + api_key: str, + output_type: int = 2, + testmode: int = 1, + dbmaski: int = 32768, + db: int = 5, + numres: int = 5, + ) -> None: + self.params = SauceNAORequest( + api_key=api_key, + output_type=output_type, + testmode=testmode, + dbmaski=dbmaski, + db=db, + numres=numres, + ) - async def _request(self, url: str): - self.params["url"] = url - try: - res = await request.get(URL, params=self.params) - except Exception: - raise RequestError("Request failed!") - data = res.json() - return data + async def _request(self, url: str) -> SauceNAOResponse: + self.params.url = url + resp = await request.get(SAUCENAO_URL, params=self.params.dict()) + return SauceNAOResponse.parse_obj(resp.json()) async def search(self, url: str) -> str: - data = await self._request(url) try: - res = data.get("results", "result") - except Exception: - return "没有相似的结果呢..." + data = await self._request(url) + except Exception as err: + raise Exception(f"处理 SauceNAO 数据失败:{str(err)}") - r = list() + r: List[SauceNAOResult] = list() for i in range(3): - data = res[i] - - sim = data["header"]["similarity"] + _data = data.results[i] + sim = _data.header.similarity if float(sim) >= 70: - _result = dict() - _result["similarity"] = sim - _result["index_name"] = data["header"]["index_name"] - _result["url"] = choice(data["data"].get("ext_urls", ["None"])) - r.append(_result) + r.append( + SauceNAOResult( + similarity=sim, + index_name=_data.header.index_name, + url=_data.data.ext_urls[0] if _data.data.ext_urls else "None", + ) + ) if not r: - return "没有相似的结果呢..." + return "SauceNAO 中没有相似的结果" - msg0 = str() + result = str() for i in r: - msg0 += ( - "\n——————————\n" - f"Similarity: {i['similarity']}\n" - f"Name: {i['index_name']}\n" - f"URL: {i['url'].replace('https://', '')}" + result += ( + MessageBuilder("\n——————————") + .text(f"相似度:{i.similarity}") + .text(f"名称:{i.index_name}") + .text(f"URL: {i.url.replace('https://', str()).replace('http://', str())}") + .done() ) - - return msg0 + return result |