🐛 修复重启命令不会触发on_shutdownhook以及可能出现的失败

LittlePaimon/plugins/bot_manager/__init__.py

@@ -3,13 +3,19 @@ import asyncio
 import random
 import sys
 from pathlib import Path
-from nonebot import on_command, get_bot
+from nonebot import on_command, get_bot, get_app
 from nonebot.permission import SUPERUSER
 from nonebot.plugin import PluginMetadata
 from nonebot.rule import to_me
 from nonebot.params import CommandArg, ArgPlainText, Arg
 from nonebot.typing import T_State
-from nonebot.adapters.onebot.v11 import Bot, Message, MessageEvent, GroupMessageEvent, ActionFailed
+from nonebot.adapters.onebot.v11 import (
+    Bot,
+    Message,
+    MessageEvent,
+    GroupMessageEvent,
+    ActionFailed,
+)
 
 from LittlePaimon.config import config
 from LittlePaimon.utils import NICKNAME, DRIVER, __version__
@@ -24,39 +30,74 @@ __plugin_meta__ = PluginMetadata(
         'author': '惜月',
         'version': '3.0',
         'priority': 99,
-    }
+    },
 )
 
-update_cmd = on_command('更新', permission=SUPERUSER, rule=to_me(), priority=1, block=True, state={
+update_cmd = on_command(
+    '更新',
+    permission=SUPERUSER,
+    rule=to_me(),
+    priority=1,
+    block=True,
+    state={
         'pm_name': 'bot_update',
         'pm_description': '从Git中更新bot，需超级用户权限',
         'pm_usage': '@bot 更新',
-    'pm_priority':    2
+        'pm_priority': 2,
-})
+    },
-check_update_cmd = on_command('检查更新', permission=SUPERUSER, rule=to_me(), priority=1, block=True, state={
+)
+check_update_cmd = on_command(
+    '检查更新',
+    permission=SUPERUSER,
+    rule=to_me(),
+    priority=1,
+    block=True,
+    state={
         'pm_name': 'bot_check_update',
         'pm_description': '从Git检查bot更新情况，需超级用户权限',
         'pm_usage': '@bot 检查更新',
-    'pm_priority':    1
+        'pm_priority': 1,
-})
+    },
-reboot_cmd = on_command('重启', permission=SUPERUSER, rule=to_me(), priority=1, block=True, state={
+)
+reboot_cmd = on_command(
+    '重启',
+    permission=SUPERUSER,
+    rule=to_me(),
+    priority=1,
+    block=True,
+    state={
         'pm_name': 'bot_restart',
         'pm_description': '执行重启操作，需超级用户权限',
         'pm_usage': '@bot 重启',
-    'pm_priority':    3
+        'pm_priority': 3,
-})
+    },
-run_cmd = on_command('cmd', permission=SUPERUSER, rule=to_me(), priority=1, block=True, state={
+)
+run_cmd = on_command(
+    'cmd',
+    permission=SUPERUSER,
+    rule=to_me(),
+    priority=1,
+    block=True,
+    state={
         'pm_name': 'bot_cmd',
         'pm_description': '运行终端命令，需超级用户权限',
         'pm_usage': '@bot cmd<命令>',
-    'pm_priority':    4
+        'pm_priority': 4,
-})
+    },
-broadcast = on_command('广播', permission=SUPERUSER, rule=to_me(), priority=1, block=True, state={
+)
+broadcast = on_command(
+    '广播',
+    permission=SUPERUSER,
+    rule=to_me(),
+    priority=1,
+    block=True,
+    state={
         'pm_name': 'broadcast',
         'pm_description': '向指定或所有群发送消息，需超级用户权限',
         'pm_usage': '@bot 广播<内容>',
-    'pm_priority':    5
+        'pm_priority': 5,
-})
+    },
+)
 
 
 @update_cmd.handle()
@@ -75,22 +116,32 @@ async def _(event: MessageEvent):
 @reboot_cmd.handle()
 async def _(bot: Bot, event: MessageEvent):
     await reboot_cmd.send(f'{NICKNAME}开始执行重启，请等待{NICKNAME}的归来', at_sender=True)
-    reboot_data = {'session_type': event.message_type,
+    reboot_data = {
-                   'session_id':   event.group_id if isinstance(event, GroupMessageEvent) else event.user_id,
+        'session_type': event.message_type,
-                   'group_card':   {}}
+        'session_id': event.group_id
+        if isinstance(event, GroupMessageEvent)
+        else event.user_id,
+        'group_card': {},
+    }
     group_list = await bot.get_group_list()
     group_id_list = [g['group_id'] for g in group_list]
     for group_id in group_id_list:
         if group_id in config.reboot_card_enable:
-            member_info = await bot.get_group_member_info(group_id=group_id, user_id=int(bot.self_id), no_cache=True)
+            member_info = await bot.get_group_member_info(
+                group_id=group_id, user_id=int(bot.self_id), no_cache=True
+            )
             reboot_data['group_card'][str(group_id)] = member_info['card']
-            await bot.set_group_card(group_id=group_id, user_id=int(bot.self_id),
+            await bot.set_group_card(
-                                     card=(member_info['card'] or member_info['nickname']) + '(重启中)')
+                group_id=group_id,
+                user_id=int(bot.self_id),
+                card=(member_info['card'] or member_info['nickname']) + '(重启中)',
+            )
             await asyncio.sleep(0.25)
     save_json(reboot_data, Path() / 'rebooting.json')
+    await get_app().router.shutdown()
     if sys.argv[0].endswith('nb'):
         sys.argv[0] = 'bot.py'
-    os.execv(sys.executable, ['python'] + sys.argv)
+    os.execv(sys.executable, [sys.executable] + sys.argv)
 
 
 @run_cmd.handle()
@@ -102,8 +153,9 @@ async def _(event: MessageEvent, state: T_State, cmd: Message = CommandArg()):
 @run_cmd.got('cmd', prompt='你输入你要运行的命令')
 async def _(event: MessageEvent, cmd: str = ArgPlainText('cmd')):
     await run_cmd.send(f'开始执行{cmd}...', at_sender=True)
-    p = await asyncio.subprocess.create_subprocess_shell(cmd, stdout=asyncio.subprocess.PIPE,
+    p = await asyncio.subprocess.create_subprocess_shell(
-                                                         stderr=asyncio.subprocess.PIPE)
+        cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+    )
     stdout, stderr = await p.communicate()
     try:
         result = (stdout or stderr).decode('utf-8')
@@ -121,14 +173,23 @@ async def _(event: MessageEvent, state: T_State, msg: Message = CommandArg()):
 
 
 @broadcast.got('groups', prompt='要广播到哪些群呢？多个群以空格隔开，或发送"全部"向所有群广播')
-async def _(event: MessageEvent, bot: Bot, msg: Message = Arg('msg'), groups: str = ArgPlainText('groups')):
+async def _(
+    event: MessageEvent,
+    bot: Bot,
+    msg: Message = Arg('msg'),
+    groups: str = ArgPlainText('groups'),
+):
     group_list = await bot.get_group_list()
     group_list = [g['group_id'] for g in group_list]
     if groups in {'全部', '所有', 'all'}:
         send_groups = group_list
     else:
         groups = groups.split(' ')
-        send_groups = [int(group) for group in groups if group.isdigit() and int(group) in group_list]
+        send_groups = [
+            int(group)
+            for group in groups
+            if group.isdigit() and int(group) in group_list
+        ]
     if not send_groups:
         await broadcast.finish('要广播的群未加入或参数不对', at_sender=True)
     else:
@@ -149,13 +210,19 @@ async def _():
     bot = get_bot()
     reboot_data = load_json(reboot_file)
     if reboot_data['session_type'] == 'group':
-        await bot.send_group_msg(group_id=reboot_data['session_id'],
+        await bot.send_group_msg(
-                                 message=f'{NICKNAME}已重启完成，当前版本为{__version__}')
+            group_id=reboot_data['session_id'],
+            message=f'{NICKNAME}已重启完成，当前版本为{__version__}',
+        )
     else:
-        await bot.send_private_msg(user_id=reboot_data['session_id'],
+        await bot.send_private_msg(
-                                   message=f'{NICKNAME}已重启完成，当前版本为{__version__}')
+            user_id=reboot_data['session_id'],
+            message=f'{NICKNAME}已重启完成，当前版本为{__version__}',
+        )
     if 'group_card' in reboot_data:
         for group_id, card_info in reboot_data['group_card'].items():
-            await bot.set_group_card(group_id=int(group_id), user_id=int(bot.self_id), card=card_info)
+            await bot.set_group_card(
+                group_id=int(group_id), user_id=int(bot.self_id), card=card_info
+            )
             await asyncio.sleep(0.25)
     reboot_file.unlink()

LittlePaimon/web/api/bot_info.py

@@ -6,7 +6,7 @@ from pathlib import Path
 
 from fastapi import APIRouter
 from fastapi.responses import JSONResponse
-from nonebot import get_bot
+from nonebot import get_bot, get_app
 from nonebot.adapters.onebot.v11 import Bot
 
 from LittlePaimon.utils import SUPERUSERS
@@ -18,114 +18,125 @@ from .utils import authentication
 route = APIRouter()
 
 
-@route.get('/get_group_list', response_class=JSONResponse, dependencies=[authentication()])
+@route.get(
+    '/get_group_list', response_class=JSONResponse, dependencies=[authentication()]
+)
 @cache(datetime.timedelta(minutes=3))
 async def get_group_list(include_all: bool = False):
     try:
         group_list = await get_bot().get_group_list()
-        group_list = [{'label': f'{group["group_name"]}({group["group_id"]})', 'value': group['group_id']} for group in group_list]
+        group_list = [
+            {
+                'label': f'{group["group_name"]}({group["group_id"]})',
+                'value': group['group_id'],
+            }
+            for group in group_list
+        ]
         if include_all:
             group_list.insert(0, {'label': '全局', 'value': 'all'})
-        return {
+        return {'status': 0, 'msg': 'ok', 'data': {'group_list': group_list}}
-            'status': 0,
-            'msg':    'ok',
-            'data':   {
-                'group_list': group_list
-            }
-        }
     except ValueError:
-        return {
+        return {'status': -100, 'msg': '获取群和好友列表失败，请确认已连接GOCQ'}
-            'status': -100,
-            'msg':    '获取群和好友列表失败，请确认已连接GOCQ'
-        }
 
 
-@route.get('/get_group_members', response_class=JSONResponse, dependencies=[authentication()])
+@route.get(
+    '/get_group_members', response_class=JSONResponse, dependencies=[authentication()]
+)
 @cache(datetime.timedelta(minutes=3))
 async def get_group_members(group_id: int):
     try:
         return await get_bot().get_group_member_list(group_id=group_id)
     except ValueError:
-        return {
+        return {'status': -100, 'msg': '获取群和好友列表失败，请确认已连接GOCQ'}
-            'status': -100,
-            'msg':    '获取群和好友列表失败，请确认已连接GOCQ'
-        }
 
 
-@route.get('/get_groups_and_members', response_class=JSONResponse, dependencies=[authentication()])
+@route.get(
+    '/get_groups_and_members',
+    response_class=JSONResponse,
+    dependencies=[authentication()],
+)
 @cache(datetime.timedelta(minutes=3))
 async def get_groups_and_members():
     result = []
     try:
         bot: Bot = get_bot()
     except ValueError:
-        return {
+        return {'status': -100, 'msg': '获取群和好友列表失败，请确认已连接GOCQ'}
-            'status': -100,
-            'msg':    '获取群和好友列表失败，请确认已连接GOCQ'
-        }
     group_list = await bot.get_group_list()
     friend_list = await bot.get_friend_list()
     for group in group_list:
         group_members = await bot.get_group_member_list(group_id=group['group_id'])
-        result.append({
+        result.append(
+            {
                 'left_label': f'{group["group_name"]}({group["group_id"]})',
                 'right_label': f'{group["group_name"]}(群{group["group_id"]})',
                 'value': f'群{group["group_id"]}',
-            'children':    [{'left_label':  f'{m["card"] or m["nickname"]}({m["user_id"]})',
+                'children': [
+                    {
+                        'left_label': f'{m["card"] or m["nickname"]}({m["user_id"]})',
                         'right_label': f'群({group["group_name"]}) - {m["card"] or m["nickname"]}({m["user_id"]})',
-                             'value':       f'群{group["group_id"]}.{m["user_id"]}'} for m in group_members if
+                        'value': f'群{group["group_id"]}.{m["user_id"]}',
-                            str(m['user_id']) != bot.self_id]
+                    }
-        })
+                    for m in group_members
+                    if str(m['user_id']) != bot.self_id
+                ],
+            }
+        )
         await asyncio.sleep(0.2)
     result = [
-        {
+        {'label': '群组', 'selectMode': 'tree', 'searchable': True, 'children': result},
-            'label':      '群组',
-            'selectMode': 'tree',
-            'searchable': True,
-            'children':   result
-        },
         {
             'label': '好友',
             'selectMode': 'list',
             'searchable': True,
-            'children':   [{
+            'children': [
+                {
                     'left_label': f'{f["nickname"]}({f["user_id"]})',
                     'right_label': f'{f["nickname"]}({f["user_id"]})',
-                'value':       f'{f["user_id"]}'
+                    'value': f'{f["user_id"]}',
-            } for f in friend_list if str(f['user_id']) != bot.self_id]
+                }
-        }]
+                for f in friend_list
+                if str(f['user_id']) != bot.self_id
+            ],
+        },
+    ]
     return result
 
 
-@route.get('/get_friend_list', response_class=JSONResponse, dependencies=[authentication()])
+@route.get(
+    '/get_friend_list', response_class=JSONResponse, dependencies=[authentication()]
+)
 @cache(datetime.timedelta(minutes=3))
 async def get_friend_list():
     try:
         bot: Bot = get_bot()
         friend_list = await bot.get_friend_list()
-        return [{'label': f'{friend["nickname"]}({friend["user_id"]})', 'value': friend['user_id']} for friend in
+        return [
-                friend_list]
+            {
-    except ValueError:
+                'label': f'{friend["nickname"]}({friend["user_id"]})',
-        return {
+                'value': friend['user_id'],
-            'status': -100,
-            'msg':    '获取群和好友列表失败，请确认已连接GOCQ'
             }
+            for friend in friend_list
+        ]
+    except ValueError:
+        return {'status': -100, 'msg': '获取群和好友列表失败，请确认已连接GOCQ'}
 
 
 @route.post('/bot_update', response_class=JSONResponse, dependencies=[authentication()])
 async def bot_update():
     result = await update()
-    return {
+    return {'status': 0, 'msg': result}
-        'status': 0,
-        'msg':    result
-    }
 
 
-@route.post('/bot_restart', response_class=JSONResponse, dependencies=[authentication()])
+@route.post(
+    '/bot_restart', response_class=JSONResponse, dependencies=[authentication()]
+)
 async def bot_restart():
-    save_json({'session_type': 'private',
+    save_json(
-               'session_id':   SUPERUSERS[0]},
+        {'session_type': 'private', 'session_id': SUPERUSERS[0]},
-              Path() / 'rebooting.json')
+        Path() / 'rebooting.json',
+    )
+    await get_app().router.shutdown()
     if sys.argv[0].endswith('nb'):
         sys.argv[0] = 'bot.py'
-    os.execv(sys.executable, ['python'] + sys.argv)
+    os.execv(sys.executable, [sys.executable] + sys.argv)