refactor: 优化代码结构和逻辑

Author: NekoAria

YetAnotherPicSearch/__init__.py

@@ -1,5 +1,6 @@
 import asyncio
 import re
+from asyncio import Lock
 from collections import defaultdict
 from contextlib import suppress
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple, Union
@@ -133,7 +134,7 @@ async def get_universal_img_url(url: str) -> str:
     final_url = re.sub(r"\?.*$", "", final_url)
     async with ClientSession(headers=DEFAULT_HEADERS) as session:
         async with session.get(final_url) as resp:
-            if resp.status == 200:
+            if resp.status < 400:
                 return final_url
     return url
 
@@ -170,28 +171,38 @@ async def send_result_message(
         msg_list = [
             msg.replace("❤️ 已收藏\n", "") if "已收藏" in msg else msg for msg in msg_list
         ]
+
     if isinstance(event, GroupMessageEvent):
         current_sending_lock = sending_lock[(event.group_id, "group")]
     else:
         current_sending_lock = sending_lock[(event.user_id, "private")]
+
     if flag := (config.forward_search_result and len(msg_list) > 1):
         try:
-            start_time = arrow.now()
-            async with current_sending_lock:
-                await send_forward_msg(bot, event, msg_list, index)
-                await asyncio.sleep(
-                    max(1 - (arrow.now() - start_time).total_seconds(), 0)
-                )
+            await send_message_with_lock(
+                bot, event, msg_list, current_sending_lock, index
+            )
         except ActionFailed:
             flag = False
     if not flag:
         for msg in msg_list:
-            start_time = arrow.now()
-            async with current_sending_lock:
-                await send_msg(bot, event, msg, index)
-                await asyncio.sleep(
-                    max(1 - (arrow.now() - start_time).total_seconds(), 0)
-                )
+            await send_message_with_lock(bot, event, [msg], current_sending_lock, index)
+
+
+async def send_message_with_lock(
+    bot: Bot,
+    event: MessageEvent,
+    msg_list: List[str],
+    current_sending_lock: Lock,
+    index: Optional[int] = None,
+) -> None:
+    start_time = arrow.now()
+    async with current_sending_lock:
+        if len(msg_list) == 1:
+            await send_msg(bot, event, msg_list[0], index)
+        else:
+            await send_forward_msg(bot, event, msg_list, index)
+        await asyncio.sleep(max(1 - (arrow.now() - start_time).total_seconds(), 0))
 
 
 async def send_msg(
@@ -200,6 +211,7 @@ async def send_msg(
     if index is not None:
         message = f"第 {index + 1} 张图片的搜索结果：\n{message}"
     message = f"{handle_reply_msg(event.message_id)}{message}"
+
     try:
         await bot.send_msg(
             user_id=event.user_id if isinstance(event, PrivateMessageEvent) else 0,

YetAnotherPicSearch/ascii2d.py

@@ -14,6 +14,7 @@ async def ascii2d_search(url: str, client: ClientSession, hide_img: bool) -> Lis
     color_res = await ascii2d_color.search(file=_file)
     if not color_res.raw:
         return ["Ascii2D 暂时无法使用"]
+
     async with ClientSession(headers=DEFAULT_HEADERS) as session:
         resp = await session.get(
             color_res.url.replace("/color/", "/bovw/"), proxy=config.proxy

YetAnotherPicSearch/cache.py

@@ -4,13 +4,15 @@
 
 from .config import config
 
+EXPIRE = config.cache_expire * 24 * 60 * 60
+
 
 def upsert_cache(cache: Cache, image_md5: str, mode: str, result: List[str]) -> None:
-    cache.set(f"{image_md5}_{mode}", result, expire=config.cache_expire * 24 * 60 * 60)
+    cache.set(f"{image_md5}_{mode}", result, expire=EXPIRE)
 
 
 def exist_in_cache(cache: Cache, image_md5: str, mode: str) -> Optional[List[str]]:
     cache_result: Optional[List[str]] = cache.get(f"{image_md5}_{mode}")
     if cache_result:
-        cache.touch(f"{image_md5}_{mode}", expire=config.cache_expire * 24 * 60 * 60)
+        cache.touch(f"{image_md5}_{mode}", expire=EXPIRE)
     return cache_result

YetAnotherPicSearch/ehentai.py

@@ -25,27 +25,32 @@
 async def ehentai_search(url: str, client: ClientSession, hide_img: bool) -> List[str]:
     ex = bool(config.exhentai_cookies)
     ehentai = EHentai(client=client)
+
     if res := await ehentai.search(url, ex=ex):
         if "Please wait a bit longer between each file search" in res.origin:
             await sleep(30 / 4)
             return await ehentai_search(url, client, hide_img)
+
         if not res.raw:
             # 如果第一次没找到，使搜索结果包含被删除的部分，并重新搜索
             async with ClientSession(headers=EHENTAI_HEADERS) as session:
                 resp = await session.get(f"{res.url}&fs_exp=on", proxy=config.proxy)
                 res = EHentaiResponse(await resp.text(), str(resp.url))
         final_res: List[str] = await search_result_filter(res, hide_img)
+
         if not res.raw and config.auto_use_ascii2d:
             final_res.append("自动使用 Ascii2D 进行搜索")
             final_res.extend(await ascii2d_search(url, client, hide_img))
         return final_res
+
     return ["EHentai 暂时无法使用"]
 
 
 async def ehentai_title_search(title: str, hide_img: bool) -> List[str]:
     title = re.sub(r"●|~| ::: |[中国翻訳]", " ", title).strip()
     url = "https://exhentai.org" if config.exhentai_cookies else "https://e-hentai.org"
     params: Dict[str, Any] = {"f_search": title}
+
     async with ClientSession(headers=EHENTAI_HEADERS) as session:
         resp = await session.get(url, proxy=config.proxy, params=params)
         if res := EHentaiResponse(await resp.text(), str(resp.url)):
@@ -56,6 +61,7 @@ async def ehentai_title_search(title: str, hide_img: bool) -> List[str]:
                 params["f_sh"] = "on"
                 resp = await session.get(url, proxy=config.proxy, params=params)
                 res = EHentaiResponse(await resp.text(), str(resp.url))
+
             # 只保留标题和搜索关键词相关度较高的结果，并排序，以此来提高准确度
             if res.raw:
                 raw_with_ratio = [
@@ -68,6 +74,7 @@ async def ehentai_title_search(title: str, hide_img: bool) -> List[str]:
                 else:
                     res.raw = [i[0] for i in raw_with_ratio]
             return await search_result_filter(res, hide_img)
+
         return ["EHentai 暂时无法使用"]
 
 
@@ -78,14 +85,17 @@ async def search_result_filter(
     _url = await shorten_url(res.url)
     if not res.raw:
         return [f"EHentai 搜索结果为空\n搜索页面：{_url}"]
+
     # 尝试过滤已删除的
     if not_expunged_res := [
         i for i in res.raw if not PyQuery(i.origin)("[id^='posted'] s")
     ]:
         res.raw = not_expunged_res
+
     # 尝试过滤无主题的杂图图集
     if not_themeless_res := [i for i in res.raw if "themeless" not in " ".join(i.tags)]:
         res.raw = not_themeless_res
+
     # 尝试过滤评分只有 1 星的
     if not_1_star_res := [
         i for i in res.raw if ("-64px" not in PyQuery(i.origin)("div.ir").attr("style"))
@@ -102,9 +112,10 @@ async def search_result_filter(
     priority["Asian Porn"] = 6
     res.raw.sort(key=lambda x: priority[x.type], reverse=True)
     for key, group in itertools.groupby(res.raw, key=lambda x: x.type):  # type: ignore
-        group_list = list(group)
-        if priority[key] > 0 and len(res.raw) != len(group_list):
-            res.raw = [i for i in res.raw if i not in group_list]
+        if priority[key] > 0:
+            group_list = list(group)
+            if len(res.raw) != len(group_list):
+                res.raw = [i for i in res.raw if i not in group_list]
 
     # 优先找汉化版；没找到就优先找原版
     if chinese_res := [

YetAnotherPicSearch/iqdb.py

@@ -14,20 +14,23 @@ async def iqdb_search(url: str, client: ClientSession, hide_img: bool) -> List[s
     res = await iqdb.search(url)
     if not res.raw:
         return ["Iqdb 暂时无法使用"]
+
     final_res: List[str] = []
     # 如果遇到搜索结果相似度低的情况，去除第一个只有提示信息的空结果
     low_acc = False
     if res.raw[0].content == "No relevant matches":
         low_acc = True
         res.raw.pop(0)
     selected_res = res.raw[0]
+
     # 优先取 danbooru 或 yande.re
     danbooru_res_list = [i for i in res.raw if i.source == "Danbooru"]
     yandere_res_list = [i for i in res.raw if i.source == "yande.re"]
     if danbooru_res_list:
         selected_res = danbooru_res_list[0]
     elif yandere_res_list:
         selected_res = yandere_res_list[0]
+
     thumbnail = await handle_img(selected_res.thumbnail, hide_img)
     source = await get_source(selected_res.url)
     if source:

YetAnotherPicSearch/utils.py

@@ -24,7 +24,7 @@ async def get_image_bytes_by_url(
     headers = {"Cookie": cookies, **DEFAULT_HEADERS} if cookies else DEFAULT_HEADERS
     async with ClientSession(headers=headers) as session:
         async with session.get(url, proxy=config.proxy) as resp:
-            if resp.status == 200 and (image_bytes := await resp.read()):
+            if resp.status < 400 and (image_bytes := await resp.read()):
                 return image_bytes
     return None
 
@@ -79,22 +79,22 @@ def handle_reply_msg(message_id: int) -> str:
 
 async def get_source(url: str) -> str:
     source = url
-    if host := URL(url).host:
+    if host := URL(source).host:
         async with ClientSession(
             headers=None if host == "danbooru.donmai.us" else DEFAULT_HEADERS
         ) as session:
-            if host in ["danbooru.donmai.us", "gelbooru.com"]:
-                async with session.get(url, proxy=config.proxy) as resp:
-                    if resp.status == 200:
-                        html = await resp.text()
-                        source = PyQuery(html)(".image-container").attr(
-                            "data-normalized-source"
-                        )
-            elif host in ["yande.re", "konachan.com"]:
-                async with session.get(url, proxy=config.proxy) as resp:
-                    if resp.status == 200:
-                        html = await resp.text()
-                        source = PyQuery(html)("#post_source").attr("value")
+            async with session.get(source, proxy=config.proxy) as resp:
+                if resp.status >= 400:
+                    return ""
+
+                html = await resp.text()
+                if host in ["danbooru.donmai.us", "gelbooru.com"]:
+                    source = PyQuery(html)(".image-container").attr(
+                        "data-normalized-source"
+                    )
+
+                elif host in ["yande.re", "konachan.com"]:
+                    source = PyQuery(html)("#post_source").attr("value")
                     if not source:
                         source = PyQuery(html)('a[href^="/pool/show/"]').text()
 
@@ -118,26 +118,30 @@ async def shorten_url(url: str) -> str:
     )
     if pid_search.search(url):
         return confuse_url(f"https://pixiv.net/i/{pid_search.search(url)[1]}")  # type: ignore
+
     uid_search = re.compile(r"pixiv.+(?:member\.php\?id=|users/)(\d+)")
     if uid_search.search(url):
         return confuse_url(f"https://pixiv.net/u/{uid_search.search(url)[1]}")  # type: ignore
+
     if URL(url).host == "danbooru.donmai.us":
         return confuse_url(url.replace("/post/show/", "/posts/"))
+
     if URL(url).host in ["exhentai.org", "e-hentai.org", "graph.baidu.com"]:
         flag = len(url) > 1024
         async with ClientSession(headers=DEFAULT_HEADERS) as session:
             if not flag:
                 resp = await session.post("https://yww.uy/shorten", json={"url": url})
-                if resp.status == 200:
+                if resp.status < 400:
                     return (await resp.json())["url"]  # type: ignore
                 else:
                     flag = True
             if flag:
                 resp = await session.post(
                     "https://www.shorturl.at/shortener.php", data={"u": url}
                 )
-                if resp.status == 200:
+                if resp.status < 400:
                     html = await resp.text()
                     final_url = PyQuery(html)("#shortenurl").attr("value")
                     return f"https://{final_url}"
+
     return confuse_url(url)

YetAnotherPicSearch/whatanime.py

@@ -18,6 +18,7 @@ async def whatanime_search(
         minutes = math.floor(time / 60)
         seconds = math.floor(time % 60)
         time_str = f"{minutes:02d}:{seconds:02d}"
+
         if res.raw[0].isAdult:
             thumbnail = await handle_img(
                 res.raw[0].cover_image,
@@ -28,12 +29,10 @@ async def whatanime_search(
                 res.raw[0].cover_image,
                 hide_img,
             )
+
         chinese_title = res.raw[0].title_chinese
         native_title = res.raw[0].title_native
 
-        def date_to_str(date: Dict[str, Any]) -> str:
-            return f"{date['year']}-{date['month']}-{date['day']}"
-
         start_date = date_to_str(res.raw[0].start_date)
         end_date = ""
         if (end_date_year := res.raw[0].end_date["year"]) and end_date_year > 0:
@@ -50,4 +49,9 @@ def date_to_str(date: Dict[str, Any]) -> str:
             f"完结：{end_date}" if end_date else "",
         ]
         return ["\n".join([i for i in res_list if i])]
+
     return ["WhatAnime 暂时无法使用"]
+
+
+def date_to_str(date: Dict[str, Any]) -> str:
+    return f"{date['year']}-{date['month']}-{date['day']}"