🎨 改进结构和代码格式

Author: NekoAria

YetAnotherPicSearch/__init__.py

@@ -1,7 +1,8 @@
 import re
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import arrow
-from nonebot.adapters.onebot.v11 import (
+from nonebot.adapters.onebot.v11 import (  # type: ignore
     ActionFailed,
     Bot,
     GroupMessageEvent,
@@ -19,13 +20,9 @@
 from .ascii2d import ascii2d_search
 from .cache import clear_expired_cache, exist_in_cache
 from .config import config
+from .result import Result
 from .saucenao import saucenao_search
 
-if config.proxy:
-    PROXY = config.proxy
-else:
-    PROXY = None
-
 
 async def _to_me(bot: Bot, event: MessageEvent) -> bool:
     msgs = event.message
@@ -52,46 +49,52 @@ async def handle_first_receive(matcher: Matcher, args: Message = CommandArg()) -
 
 
 async def image_search(
-    url: str, mode: str, purge: bool, proxy: str, hide_img: bool = config.hide_img
-) -> list[str]:
+    url: str,
+    mode: str,
+    purge: bool,
+    proxy: Optional[str],
+    hide_img: bool = config.hide_img,
+) -> Union[List[str], Any]:
     db = TinyDB(
         "cache.json",
-        storage=CachingMiddleware(JSONStorage),
+        storage=CachingMiddleware(JSONStorage),  # type: ignore
         encoding="utf-8",
         sort_keys=True,
         indent=4,
         ensure_ascii=False,
     )
-    image_md5 = re.search("[A-F0-9]{32}", url)[0]
-    result = await exist_in_cache(db, image_md5, mode)
-    cached = bool(result)
-    if purge or not cached:
-        result = {}
+    image_md5 = re.search("[A-F0-9]{32}", url)[0]  # type: ignore
+    _result = await exist_in_cache(db, image_md5, mode)
+    cached = bool(_result)
+    if purge or not _result:
+        result_dict: Dict[str, Any] = {}
         if mode == "a2d":
-            result["ascii2d"] = await ascii2d_search(url, proxy, hide_img)
+            result_dict["ascii2d"] = await ascii2d_search(url, proxy, hide_img)
         else:
-            result["saucenao"] = await saucenao_search(url, mode, proxy, hide_img)
-        result["mode"] = mode
-        result["image_md5"] = image_md5
-        result["update_at"] = arrow.now().for_json()
-        db.upsert(result, (Query().image_md5 == image_md5) & (Query().mode == mode))
-        db.insert(result)
+            result_dict["saucenao"] = await saucenao_search(url, mode, proxy, hide_img)
+        result_dict["mode"] = mode
+        result_dict["image_md5"] = image_md5
+        result_dict["update_at"] = arrow.now().for_json()
+        _result = Result(result_dict)
+        db.upsert(
+            _result.__dict__, (Query().image_md5 == image_md5) & (Query().mode == mode)
+        )
     await clear_expired_cache(db)
     db.close()
     if mode == "a2d":
-        final_res = result["ascii2d"]
+        final_res = _result.ascii2d
     else:
-        final_res = result["saucenao"]
+        final_res = _result.saucenao
     if cached and not purge:
-        final_res = [f"[缓存] {i}" for i in final_res]
+        return [f"[缓存] {i}" for i in final_res]  # type: ignore
     return final_res
 
 
-async def get_image_urls(msg: Message) -> list[str]:
+async def get_image_urls(msg: Message) -> List[str]:
     return [i.data["url"] for i in msg if i.type == "image" and i.data.get("url")]
 
 
-async def get_args(msg: Message) -> (str, bool):
+async def get_args(msg: Message) -> Tuple[str, bool]:
     mode = "all"
     plain_text = msg.extract_plain_text()
     args = ["a2d", "pixiv", "danbooru", "doujin", "anime"]
@@ -112,7 +115,7 @@ async def handle_image_search(bot: Bot, event: MessageEvent) -> None:
         await IMAGE_SEARCH_MODE.reject()
     mode, purge = await get_args(event.message)
     for i in image_urls:
-        msg_list = await image_search(i, mode, purge, PROXY)
+        msg_list = await image_search(i, mode, purge, config.proxy)
         if isinstance(event, PrivateMessageEvent):
             for msg in msg_list:
                 await bot.send_private_msg(user_id=event.user_id, message=msg)

YetAnotherPicSearch/ascii2d.py

@@ -1,10 +1,12 @@
-from PicImageSearch import Ascii2D, Network
-from PicImageSearch.Utils import Ascii2DResponse
+from typing import List, Optional
+
+from PicImageSearch import Ascii2D, Network  # type: ignore
+from PicImageSearch.Utils import Ascii2DResponse  # type: ignore
 
 from .utils import handle_img, shorten_url
 
 
-async def ascii2d_search(url: str, proxy: str, hide_img: bool) -> list[str]:
+async def ascii2d_search(url: str, proxy: Optional[str], hide_img: bool) -> List[str]:
     async with Network(proxies=proxy) as client:
         ascii2d_color = Ascii2D(client=client)
         ascii2d_bovw = Ascii2D(bovw=True, client=client)
@@ -15,11 +17,14 @@ async def get_final_res(res: Ascii2DResponse) -> str:
             if not res.raw[0].url:
                 res.raw[0] = res.raw[1]
             thumbnail = await handle_img(res.raw[0].thumbnail, proxy, hide_img)
+            _url = ""
+            if res.raw[0]:
+                _url = await shorten_url(res.raw[0].url)
             res_list = [
                 f"{thumbnail}",
                 f"{res.raw[0].title}" if res.raw[0].title else "",
                 f"Author：{res.raw[0].author}" if res.raw[0].author else "",
-                f"{await shorten_url(res.raw[0].url)}" if res.raw[0].url else "",
+                _url,
             ]
             return "\n".join([i for i in res_list if i != ""])
 

YetAnotherPicSearch/cache.py

@@ -1,19 +1,25 @@
+from typing import List, Optional
+
 import arrow
 from tinydb import Query, TinyDB
 from tinydb.operations import set
+from tinydb.table import Document
 
 from .config import config
+from .result import Result
 
 
-async def exist_in_cache(db: TinyDB, image_md5: str, mode: str) -> dict:
-    cache_result = db.search((Query().image_md5 == image_md5) & (Query().mode == mode))
+async def exist_in_cache(db: TinyDB, image_md5: str, mode: str) -> Optional[Result]:
+    result: Optional[Result] = None
+    cache_result: List[Document] = db.search(
+        (Query().image_md5 == image_md5) & (Query().mode == mode)
+    )
     if cache_result:
         db.update(
-            set("update_at", arrow.now().for_json()), Query().image_md5 == image_md5
+            set("update_at", arrow.now().for_json()), Query().image_md5 == image_md5  # type: ignore
         )
-        return cache_result[-1]
-    else:
-        return {}
+        result = Result(cache_result[-1])
+    return result
 
 
 async def clear_expired_cache(db: TinyDB) -> None:

YetAnotherPicSearch/config.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from nonebot import get_driver
 from pydantic import BaseSettings
 
@@ -8,7 +10,7 @@ class Config(BaseSettings):
     saucenao_low_acc: int = 60  # saucenao 相似度低于这个百分比将被认定为相似度过低
     use_ascii2d_when_low_acc: bool = True  # 是否在 saucenao 相似度过低时自动使用 ascii2d
     group_forward_search_result: bool = True  # 若结果消息有多条，采用合并转发方式发送搜图结果（私聊不生效）
-    proxy: str = ""  # 大部分请求所使用的代理，支持 http(s):// 和 socks://
+    proxy: Optional[str] = None  # 大部分请求所使用的代理: http(s)://
     cache_expire: int = 7  # 搜图结果缓存过期时间（天）
     saucenao_api_key: str = ""  # saucenao APIKEY，必填，否则无法使用 saucenao 搜图
 

YetAnotherPicSearch/result.py

@@ -0,0 +1,12 @@
+from typing import Any, Dict, List, Optional
+
+
+class Result:
+    def __init__(self, mapping: Dict[str, Any]):
+        self.ascii2d: Optional[List[str]] = []
+        self.saucenao: Optional[List[str]] = []
+        self.image_md5: Optional[str] = None
+        self.update_at: Optional[str] = None
+        self.mode: str = "all"
+        for key, value in mapping.items():
+            setattr(self, key, value)

YetAnotherPicSearch/saucenao.py

@@ -1,13 +1,16 @@
 import re
+from typing import List, Optional
 
-from PicImageSearch import Network, SauceNAO
+from PicImageSearch import Network, SauceNAO  # type: ignore
 
 from .ascii2d import ascii2d_search
 from .config import config
 from .utils import get_source, handle_img, shorten_url
 
 
-async def saucenao_search(url: str, mode: str, proxy: str, hide_img: bool) -> list[str]:
+async def saucenao_search(
+    url: str, mode: str, proxy: Optional[str], hide_img: bool
+) -> List[str]:
     saucenao_db = {"all": 999, "pixiv": 5, "danbooru": 9, "anime": 21, "doujin": 18}
     async with Network(proxies=proxy) as client:
         saucenao = SauceNAO(
@@ -31,7 +34,7 @@ async def saucenao_search(url: str, mode: str, proxy: str, hide_img: bool) -> li
                 if len(pixiv_res_list) > 1:
                     selected_res = min(
                         pixiv_res_list,
-                        key=lambda x: int(re.search(r"\d+", x.url).group()),
+                        key=lambda x: int(re.search(r"\d+", x.url).group()),  # type: ignore
                     )
             # 如果地址有多个，优先取 danbooru
             elif ext_urls and len(ext_urls) > 1:
@@ -48,12 +51,13 @@ async def saucenao_search(url: str, mode: str, proxy: str, hide_img: bool) -> li
                 "data"
             ].get("jp_name"):
                 selected_res.title = selected_res.origin["data"]["jp_name"]
+            _url = await shorten_url(selected_res.url)
             res_list = [
                 f"SauceNAO（{selected_res.similarity}%）",
                 f"{thumbnail}",
                 f"{selected_res.title}",
                 f"Author：{selected_res.author}" if selected_res.author else "",
-                f"{await shorten_url(selected_res.url)}",
+                _url,
                 f"Source：{source}" if source else "",
             ]
             final_res.append("\n".join([i for i in res_list if i != ""]))

YetAnotherPicSearch/utils.py

@@ -1,20 +1,21 @@
 import base64
 import re
 from io import BytesIO
+from typing import Optional
 
 import httpx
-from pyquery import PyQuery
+from pyquery import PyQuery  # type: ignore
 
 
 # 将图片转化为 base64
-async def get_pic_base64_by_url(url: str, proxy: str) -> str:
-    async with httpx.AsyncClient(proxies=proxy) as client:
+async def get_pic_base64_by_url(url: str, proxy: Optional[str]) -> str:
+    async with httpx.AsyncClient(proxies=proxy) as client:  # type: ignore
         r = await client.get(url)
         image_buffer = BytesIO(r.content)
         return str(base64.b64encode(image_buffer.getvalue()), encoding="utf-8")
 
 
-async def handle_img(url: str, proxy: str, hide_img: bool) -> str:
+async def handle_img(url: str, proxy: Optional[str], hide_img: bool) -> str:
     if hide_img:
         return ""
     img_base64 = await get_pic_base64_by_url(url, proxy)
@@ -23,9 +24,9 @@ async def handle_img(url: str, proxy: str, hide_img: bool) -> str:
     return f"图片下载失败: {url}"
 
 
-async def get_source(url: str, proxy: str) -> str:
+async def get_source(url: str, proxy: Optional[str]) -> str:
     source = ""
-    async with httpx.AsyncClient(proxies=proxy) as client:
+    async with httpx.AsyncClient(proxies=proxy) as client:  # type: ignore
         if httpx.URL(url).host == "danbooru.donmai.us":
             source = PyQuery((await client.get(url)).content)(".image-container").attr(
                 "data-normalized-source"
@@ -46,7 +47,7 @@ async def get_source(url: str, proxy: str) -> str:
 async def shorten_url(url: str) -> str:
     pid_search = re.compile(r"pixiv.+(?:illust_id=|artworks/)([0-9]+)")
     if pid_search.search(url):
-        return f"https://pixiv.net/i/{pid_search.search(url).group(1)}"
+        return f"https://pixiv.net/i/{pid_search.search(url).group(1)}"  # type: ignore
     if httpx.URL(url).host == "danbooru.donmai.us":
         return url.replace("/post/show/", "/posts/")
     return url