1 files changed, 53 insertions, 62 deletions
diff --git a/ATRI/plugins/setu/data_source.py b/ATRI/plugins/setu/data_source.py
index 79a36e0..59f511a 100644
--- a/ATRI/plugins/setu/data_source.py
+++ b/ATRI/plugins/setu/data_source.py
@@ -1,81 +1,72 @@
-import asyncio
-from nonebot.adapters.onebot.v11 import Bot, MessageSegment
+from typing import Tuple
+from nonebot.adapters.onebot.v11 import MessageSegment
 
 from ATRI import conf
 from ATRI.utils import request
+from ATRI.exceptions import RequestError
+
+from .models import SetuInfo
 from .nsfw_checker import detect_image, init_model
 
 
-LOLICON_URL = "https://api.lolicon.app/setu/v2"
-DEFAULT_SETU = (
-    "https://i.pixiv.cat/img-original/img/2021/02/28/22/44/49/88124144_p0.jpg"
-)
+_LOLICON_URL = "https://api.lolicon.app/setu/v2"
 
 
 class Setu:
-    @staticmethod
-    def _use_proxy(url: str) -> str:
-        if conf.Setu.reverse_proxy:
-            return url.replace("i.pixiv.cat", conf.Setu.reverse_proxy_domain)
-        else:
-            return url
+    def __init__(self, url: str):
+        self.url = url
 
     @classmethod
-    async def random_setu(cls) -> tuple:
-        """
-        随机涩图.
-        """
-        res = await request.get(LOLICON_URL)
-        data: dict = res.json()
-        temp_data: dict = data.get("data", list())
-        if not temp_data:
-            return "涩批爬", None
+    async def new(cls, tag: str = str()) -> Tuple[MessageSegment, SetuInfo]:
+        """new 一个涩图
 
-        data: dict = temp_data[0]
-        title = data.get("title", "木陰のねこ")
-        p_id = data.get("pid", 88124144)
-        url: str = data["urls"].get("original", "ignore")
+        Args:
+            tag (str, optional): 附加 tag, 默认无
 
-        setu = MessageSegment.image(cls._use_proxy(url), timeout=114514)
-        repo = f"Title: {title}\nPid: {p_id}"
-        return repo, setu
+        Raises:
+            RequestError: 涩图请求失败
 
-    @classmethod
-    async def tag_setu(cls, tag: str) -> tuple:
-        """
-        指定tag涩图.
-        """
-        url = LOLICON_URL + f"?tag={tag}"
-        res = await request.get(url)
-        data: dict = res.json()
-
-        temp_data: dict = data.get("data", list())
-        if not temp_data:
-            return f"没有 {tag} 的涩图呢...", None
-
-        data = temp_data[0]
-        title = data.get("title", "木陰のねこ")
-        p_id = data.get("pid", 88124144)
-        url = data["urls"].get(
-            "original",
-            cls._use_proxy(DEFAULT_SETU),
-        )
-        setu = MessageSegment.image(url, timeout=114514)
-        repo = f"Title: {title}\nPid: {p_id}"
-        return repo, setu
-
-    @staticmethod
-    async def detecter(url: str, file_size: int) -> float:
-        """
-        涩值检测.
+        Returns:
+            Tuple[MessageSegment, dict]: 涩图本体, 涩图信息
         """
-        data = await detect_image(url, file_size)
-        return data
+        url = _LOLICON_URL
+        if tag:
+            url = _LOLICON_URL + f"?tag={tag}"
+        try:
+            req = await request.get(url)
+        except Exception:
+            raise RequestError("setu: 请求失败")
+
+        data = req.json()
+        cache_data = data.get("data")
+        if not cache_data:
+            raise RequestError("今天不可以涩")
+
+        data = cache_data[0]
+        title = data["title"]
+        pid = data["pid"]
+        setu = data["urls"].get("original", "ignore")
 
-    @staticmethod
-    async def async_recall(bot: Bot, event_id):
-        await asyncio.sleep(30)
-        await bot.delete_msg(message_id=event_id)
+        if conf.Setu.reverse_proxy:
+            setu = MessageSegment.image(
+                file=setu.replace("i.pixiv.cat", conf.Setu.reverse_proxy_domain),
+                timeout=114514,
+            )
+
+        setu_data = SetuInfo(title=title, pid=pid)
+
+        return setu, setu_data
+
+    async def detecter(self, max_size: int) -> float:
+        """图片涩值检测
+
+        Args:
+            max_size (int): 检测文件大小限制
+
+        Returns:
+            float: 百分比涩值
+        """
+        return await detect_image(self.url, max_size)
 
 
 from ATRI import driver