nonebot-plugin-shitbot 0.4.2__tar.gz → 0.5.0__tar.gz

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nonebot-plugin-shitbot
-Version: 0.4.2
+Version: 0.5.0
 Summary: NoneBot shitpost plugin
 Requires-Python: >=3.10
 License-File: LICENSE

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/aux.py

@@ -6,8 +6,6 @@ from typing import Any
 import httpx
 from nonebot.log import logger
 
-from .config import config
-
 
 def validate_schema(data: Any, schema: Any) -> bool:
     if schema is Any:
@@ -36,7 +34,7 @@ def validate_schema(data: Any, schema: Any) -> bool:
     return False
 
 
-async def rm_path(path: Path):
+def rm_path(path: Path):
     if not path.exists():
         return
     if path.is_dir():
@@ -45,12 +43,6 @@ async def rm_path(path: Path):
         path.unlink()
 
 
-async def rm_cache(group_id: str, user_id: str, pid: str):
-    cache_path = config.cache / group_id / user_id / pid
-    await rm_path(cache_path)
-    logger.info(f"已删除缓存: {cache_path}")
-
-
 def get_file_sha256(file_path: Path | str):
     file = Path(file_path) if isinstance(file_path, str) else file_path
     if not file.is_file():

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/command.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import asyncio
+from functools import wraps
 from typing import TYPE_CHECKING
 
 from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent
 
+from .aux import rm_path
+from .config import config
 from .msgutils import send_msg
 from .parser import BotArgParser
 from .permissions import permissions
@@ -12,6 +16,40 @@ if TYPE_CHECKING:
     from .session import BotSession
 
 
+def _cancellable_run(func):
+    @wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        task = asyncio.current_task()
+        if task is not None:
+            self._run_tasks.add(task)
+        try:
+            return await func(self, *args, **kwargs)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            if task is not None:
+                self._run_tasks.discard(task)
+
+    return wrapper
+
+
+def _cancellable_roger(func):
+    @wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        task = asyncio.current_task()
+        if task is not None:
+            self._roger_tasks.add(task)
+        try:
+            return await func(self, *args, **kwargs)
+        except asyncio.CancelledError:
+            pass
+        finally:
+            if task is not None:
+                self._roger_tasks.discard(task)
+
+    return wrapper
+
+
 # Base class of all the commands
 class BotCommand:
     # class fields:
@@ -27,6 +65,13 @@ class BotCommand:
     _sentinel = object()
     _name = "otherwise"
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "run" in cls.__dict__:
+            cls.run = _cancellable_run(cls.run)
+        if "roger" in cls.__dict__:
+            cls.roger = _cancellable_roger(cls.roger)
+
     # Use make() instead of this
     def __init__(self, bot: Bot, session: BotSession, *, _pid: int, _internal=None):
         if _internal is not self._sentinel:
@@ -36,6 +81,10 @@ class BotCommand:
         self._argv = None
         self._parser = self._init_parser()
         self._pid = _pid
+        self._run_tasks: set[asyncio.Task] = set()
+        self._roger_tasks: set[asyncio.Task] = set()
+        self._inner_tasks: set[asyncio.Task] = set()
+        self._inner_cmds: list[BotCommand] = []
         session.commands[_pid] = self
 
     @classmethod
@@ -72,6 +121,14 @@ class BotCommand:
     def pid(self):
         return self._pid
 
+    @property
+    def run_tasks(self):
+        return self._run_tasks
+
+    @property
+    def roger_tasks(self):
+        return self._roger_tasks
+
     def _init_parser(self):
         return BotArgParser()
 
@@ -131,6 +188,20 @@ class BotCommand:
             entry_name, self.session.group_id, self.session.user_id
         )
 
+    def _create_task(self, coro) -> asyncio.Task:
+        task = asyncio.create_task(coro)
+        self._inner_tasks.add(task)
+        task.add_done_callback(self._inner_tasks.discard)
+        return task
+
+    def _rm_cache(self):
+        if not self.session:
+            return
+        cache_path = (
+            config.cache / self.session.group_id / self.session.user_id / str(self._pid)
+        )
+        rm_path(cache_path)
+
     async def roger(self, event: MessageEvent):
         pass
 
@@ -149,6 +220,16 @@ class BotCommand:
     def unlock(self):
         if self._session is None:
             return
+        for task in self._inner_tasks:
+            if not task.done():
+                task.cancel()
+        for cmd in self._inner_cmds:
+            cmd.unlock()
+        self._rm_cache()
         self._session.commands.pop(self._pid, None)
         self._session.release()
         self._session = None
+
+
+BotCommand.run = _cancellable_run(BotCommand.run)
+BotCommand.roger = _cancellable_roger(BotCommand.roger)

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/commands/convert_cmd.py

@@ -11,7 +11,7 @@ import httpx
 from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent, MessageSegment
 from nonebot.log import logger
 
-from ..aux import rm_cache, stuff_download
+from ..aux import stuff_download
 from ..command import BotCommand
 from ..config import config
 from ..msgutils import dump_message, get_message_with_reply, get_multimedias_url
@@ -96,6 +96,7 @@ class BotCommandConvert(BotCommand):
             return ""
         png_path = png_dir / f"{uuid.uuid4().hex}.png"
         import magic
+
         mime = magic.from_file(download_path, mime=True)
         logger.info(
             f"检测文件类型: {mime}, 执行图片转 PNG: {download_path} -> {png_path}"
@@ -312,6 +313,8 @@ class BotCommandConvert(BotCommand):
                         client,
                         downloads_dir,
                     )
+            except asyncio.CancelledError:
+                pass
             except Exception as e:
                 logger.exception(f"urls_to_downloads 管道异常退出: {e}")
 
@@ -320,6 +323,8 @@ class BotCommandConvert(BotCommand):
                 await prod_cons(
                     self._downloads, self._pngs, self._download_to_png, pngs_dir
                 )
+            except asyncio.CancelledError:
+                pass
             except Exception as e:
                 logger.exception(f"downloads_to_pngs 管道异常退出: {e}")
 
@@ -334,12 +339,14 @@ class BotCommandConvert(BotCommand):
                         await prod_cons(
                             self._pngs, self._outputs, self._png_to_frame, outputs_dir
                         )
+            except asyncio.CancelledError:
+                pass
             except Exception as e:
                 logger.exception(f"pngs_to_outputs 管道异常退出: {e}")
 
-        asyncio.create_task(_urls_to_downloads())
-        asyncio.create_task(_downloads_to_pngs())
-        asyncio.create_task(_pngs_to_outputs())
+        self._create_task(_urls_to_downloads())
+        self._create_task(_downloads_to_pngs())
+        self._create_task(_pngs_to_outputs())
 
         logger.info(f"用户 {self.session.user_id} 开始了图片收集")
         await self.send_msg(
@@ -363,7 +370,6 @@ class BotCommandConvert(BotCommand):
         await self.send_msg("转换完毕。")
 
         await self._send_outputs()
-        await rm_cache(self.session.group_id, self.session.user_id, str(self._pid))
         self.unlock()
         return
 

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/commands/help_cmd.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import re
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -34,6 +35,14 @@ class BotCommandHelp(BotCommand):
         parser.add_subparser("autoreply")
         return parser
 
+    def unlock(self):
+        for task in self._inner_tasks:
+            if not task.done():
+                task.cancel()
+        for cmd in self._inner_cmds:
+            cmd.unlock()
+        super().unlock()
+
     async def run(self, args: Message):
         if not self.session:
             return
@@ -88,7 +97,11 @@ class BotCommandHelp(BotCommand):
         while _pid in pids:
             _pid -= 1
         md2pic = BotCommandMd2pic.make(self.bot, self.session, _pid=_pid)
+        if md2pic is None:
+            return
         tip = "-c\n" + tip
         msg = Message(tip)
-        await md2pic.run(msg)  # type: ignore
+        inner_task = self._create_task(md2pic.run(msg))
+        self._inner_cmds.append(md2pic)
+        await asyncio.gather(inner_task, return_exceptions=True)
         self.unlock()

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/commands/session_cmd.py

@@ -23,8 +23,10 @@ class BotCommandSession(BotCommand):
         parser.set_rule(max=0, need_subcmd=True)
         switch = parser.add_subparser("switch")
         info = parser.add_subparser("info")
+        stop = parser.add_subparser("stop")
         switch.set_rule(min=1, max=1, types=[int])
         info.set_rule(max=0)
+        stop.set_rule(max=0)
         return parser
 
     async def run(self, args: Message):
@@ -76,4 +78,19 @@ class BotCommandSession(BotCommand):
             tip += f"正在运行的命令: \n{json.dumps(cmd_info, indent=2, ensure_ascii=False)}"
             await self.send_msg(tip)
 
+        if subcmd == "stop":
+            curpid = self.session.curpid
+            target = self.session.commands.get(curpid)
+            if target is None:
+                await self.send_msg(f"pid {curpid} 没有正在运行的命令")
+            else:
+                for task in target.run_tasks:
+                    if not task.done():
+                        task.cancel()
+                for task in target.roger_tasks:
+                    if not task.done():
+                        task.cancel()
+                target.unlock()
+                await self.send_msg(f"已停止命令 {target.name} (pid={curpid})")
+
         self.unlock()

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/handlers.py

@@ -1,5 +1,3 @@
-import asyncio
-
 import yaml
 from nonebot import get_driver, on_command, on_message
 from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent
@@ -89,7 +87,7 @@ driver = get_driver()
 
 @driver.on_startup
 async def startup():
-    await rm_path(config.cache)
+    rm_path(config.cache)
     config.cache.mkdir(parents=True, exist_ok=True)
     logger.info("nonebot-plugin-shitbot 已加载")
 
@@ -110,11 +108,22 @@ async def bot_connect(bot: Bot):
 
 @driver.on_shutdown
 async def shutdown():
-    tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-    for task in tasks:
-        task.cancel()
-
-    await asyncio.gather(*tasks, return_exceptions=True)
+    commands = [
+        cmd
+        for group_objs in BotSession._objs.values()
+        for session in group_objs.values()
+        for pid, cmd in session.commands.items()
+        if pid >= 0
+    ]
+    for command in commands:
+        for task in command.run_tasks:
+            if not task.done():
+                task.cancel()
+        for task in command.roger_tasks:
+            if not task.done():
+                task.cancel()
+        command.unlock()
+    rm_path(config.cache)
 
 
 # ===Commands handlers=== #

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot/help/session.md

@@ -14,6 +14,7 @@
 |--------|------|
 | `switch <pid>` | 将前台切换到指定 pid |
 | `info` | 查看当前会话信息 |
+| `stop` | 停止当前会话的前台 pid 中运行的命令 |
 
 ## 说明
 
@@ -35,4 +36,5 @@
 ```bash
 /session switch 3    将前台切换到 pid=3
 /session info        查看所有运行中的命令以及前台对应的 pid
+/session stop        停止当前会话的前台 pid 中运行的命令
 ```

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/nonebot_plugin_shitbot.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nonebot-plugin-shitbot
-Version: 0.4.2
+Version: 0.5.0
 Summary: NoneBot shitpost plugin
 Requires-Python: >=3.10
 License-File: LICENSE

{nonebot_plugin_shitbot-0.4.2 → nonebot_plugin_shitbot-0.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nonebot-plugin-shitbot"
-version = "0.4.2"
+version = "0.5.0"
 description = "NoneBot shitpost plugin"
 requires-python = ">=3.10"
 dependencies = [