update code

plugins/pix_gallery/pix.py

@@ -1,17 +1,30 @@
-from nonebot import on_command
-from nonebot.adapters.cqhttp.event import GroupMessageEvent
 from nonebot.adapters.cqhttp.message import Message
 from utils.utils import get_message_text, is_number
+from configs.config import PIX_OMEGA_PIXIV_RATIO, WITHDRAW_PIX_TIME
+from models.omega_pixiv_illusts import OmegaPixivIllusts
 from utils.message_builder import image
 from services.log import logger
-from nonebot.adapters.cqhttp import Bot, MessageEvent
+from nonebot.adapters.cqhttp import (
+    Bot,
+    MessageEvent,
+    GroupMessageEvent,
+    PrivateMessageEvent,
+)
+from utils.static_data import withdraw_message_id_manager
 from nonebot.typing import T_State
 from .data_source import get_image
 from models.pixiv import Pixiv
+from nonebot import on_command
 import random
 
 
-pix = on_command("pix", aliases={"PIX"}, priority=5, block=True)
+pix = on_command("pix", aliases={"PIX", "Pix"}, priority=5, block=True)
+
+
+PIX_RATIO = PIX_OMEGA_PIXIV_RATIO[0] / (
+    PIX_OMEGA_PIXIV_RATIO[0] + PIX_OMEGA_PIXIV_RATIO[1]
+)
+OMEGA_RATIO = 1 - PIX_RATIO
 
 
 @pix.handle()
@@ -24,40 +37,97 @@ async def _(bot: Bot, event: MessageEvent, state: T_State):
         pid = keyword.replace("pid", "").replace(":", "")
         if not is_number(pid):
             await pix.finish("PID必须是数字...", at_sender=True)
-        all_image = await Pixiv.query_images(pid=int(pid))
+        isr18 = 0 if isinstance(event, GroupMessageEvent) else 2
+        nsfw_tag = 0 if isinstance(event, GroupMessageEvent) else None
+        all_image = await Pixiv.query_images(pid=int(pid), r18=isr18)
+        if not all_image:
+            all_image = await OmegaPixivIllusts.query_images(
+                pid=int(pid), nsfw_tag=nsfw_tag
+            )
     else:
         x = keyword.split()
+        if '-s' in x:
+            x.remove('-s')
+            nsfw_tag = 1
+        elif '-r' in x:
+            x.remove('-r')
+            nsfw_tag = 2
+        else:
+            nsfw_tag = 0
+        if nsfw_tag != 0 and str(event.user_id) not in bot.config.superusers:
+            await pix.finish('你不能看这些噢，这些都是是留给管理员看的...')
         if len(x) > 1:
             if is_number(x[-1]):
                 num = int(x[-1])
-                if num > 11:
-                    num = random.randint(1, 10)
-                    await pix.send(f"太贪心了，就给你发 {num}张 好了")
+                if num > 10:
+                    if str(event.user_id) not in bot.config.superusers or (
+                        str(event.user_id) in bot.config.superusers and num > 30
+                    ):
+                        num = random.randint(1, 10)
+                        await pix.send(f"太贪心了，就给你发 {num}张 好了")
                 x = x[:-1]
                 keyword = " ".join(x)
-        all_image = await Pixiv.query_images(x)
+        pix_num = int(num * PIX_RATIO) + 15 if PIX_RATIO != 0 else 0
+        omega_num = num - pix_num + 15
+        tmp = await Pixiv.query_images(
+            x, r18=1 if nsfw_tag == 2 else 0, num=pix_num
+        ) + await OmegaPixivIllusts.query_images(x, nsfw_tag=nsfw_tag, num=omega_num)
+        tmp_ = []
+        all_image = []
+        for x in tmp:
+            if x.pid not in tmp_:
+                all_image.append(x)
+                tmp_.append(x.pid)
     if not all_image:
         await pix.finish(f"未在图库中找到与 {keyword} 相关Tag/UID/PID的图片...", at_sender=True)
     for _ in range(num):
+        img_url = None
+        author = None
         if not all_image:
             await pix.finish("坏了...发完了，没图了...")
         img = random.choice(all_image)
         all_image.remove(img)
-        img_url = img.img_url
+        if isinstance(img, OmegaPixivIllusts):
+            img_url = img.url
+            author = img.uname
+            print(img.nsfw_tag)
+        elif isinstance(img, Pixiv):
+            img_url = img.img_url
+            author = img.author
         pid = img.pid
         title = img.title
-        author = img.author
         uid = img.uid
         # tags = img.tags
-        await pix.send(
-            Message(
-                f"title：{title}\n"
-                f"author：{author}\n"
-                f"PID：{pid}\nUID：{uid}\n"
-                f"{image(await get_image(img_url, event.user_id), 'temp')}"
+        _img = await get_image(img_url, event.user_id)
+        if _img:
+            msg_id = await pix.send(
+                Message(
+                    f"title：{title}\n"
+                    f"author：{author}\n"
+                    f"PID：{pid}\nUID：{uid}\n"
+                    f"{image(_img, 'temp')}"
+                )
+            )
+            logger.info(
+                f"(USER {event.user_id}, GROUP {event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
+                f" 查看PIX图库PID: {pid}"
+            )
+        else:
+            msg_id = await pix.send(f"下载图片似乎出了一点问题，PID：{pid}")
+            logger.info(
+                f"(USER {event.user_id}, GROUP {event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
+                f" 查看PIX图库PID: {pid}，下载图片出错"
+            )
+        withdraw_message(event, msg_id["message_id"])
+
+
+def withdraw_message(event: MessageEvent, id_: int):
+    if WITHDRAW_PIX_TIME[0]:
+        if (
+            (WITHDRAW_PIX_TIME[1] == 0 and isinstance(event, PrivateMessageEvent))
+            or (WITHDRAW_PIX_TIME[1] == 1 and isinstance(event, GroupMessageEvent))
+            or WITHDRAW_PIX_TIME[1] == 2
+        ):
+            withdraw_message_id_manager["message_id"].append(
+                (id_, WITHDRAW_PIX_TIME[0])
             )
-        )
-        logger.info(
-            f"(USER {event.user_id}, GROUP {event.group_id if isinstance(event, GroupMessageEvent) else 'private'})"
-            f" 查看PIX图库PID: {pid}"
-        )