ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2024-11-28 03:15:24 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update __init__.py

Browse Source 
修复CommandSession.pause和CommandSession.finish在多线程环境下可能运行出错的bug
This commit is contained in:
Mix 2020-02-22 10:03:50 +08:00 committed by GitHub
parent 960dfd95df
commit bdb64569e7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 62 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

114
nonebot/command/__init__.py
View File

@ -3,10 +3,9 @@ import re
import shlex
import warnings
from datetime import datetime
from typing import (
Tuple, Union, Callable, Iterable, Any, Optional, List, Dict,
Awaitable
)
from functools import partial
from typing import (Tuple, Union, Callable, Iterable, Any, Optional, List,
Dict, Awaitable)
from nonebot import NoneBot, permission as perm
from nonebot.command.argfilter import ValidateError
@ -14,10 +13,8 @@ from nonebot.helpers import context_id, send, render_expression
from nonebot.log import logger
from nonebot.message import Message
from nonebot.session import BaseSession
from nonebot.typing import (
Context_T, CommandName_T, CommandArgs_T, Message_T, State_T,
Filter_T
)
from nonebot.typing import (Context_T, CommandName_T, CommandArgs_T, Message_T,
State_T, Filter_T)
# key: one segment of command name
# value: subtree or a leaf Command object
@ -35,18 +32,11 @@ CommandHandler_T = Callable[['CommandSession'], Any]
class Command:
__slots__ = ('name', 'func',
'permission',
'only_to_me',
'privileged',
__slots__ = ('name', 'func', 'permission', 'only_to_me', 'privileged',
'args_parser_func')
def __init__(self, *,
name: CommandName_T,
func: CommandHandler_T,
permission: int,
only_to_me: bool,
privileged: bool):
def __init__(self, *, name: CommandName_T, func: CommandHandler_T,
permission: int, only_to_me: bool, privileged: bool):
self.name = name
self.func = func
self.permission = permission
@ -54,7 +44,9 @@ class Command:
self.privileged = privileged
self.args_parser_func: Optional[CommandHandler_T] = None
async def run(self, session, *,
async def run(self,
session,
*,
check_perm: bool = True,
dry: bool = False) -> bool:
"""
@ -92,15 +84,16 @@ class Command:
if session.state['__validation_failure_num'] >= \
config.MAX_VALIDATION_FAILURES:
# noinspection PyProtectedMember
session.finish(render_expression(
config.TOO_MANY_VALIDATION_FAILURES_EXPRESSION
), **session._current_send_kwargs)
session.finish(
render_expression(
config.
TOO_MANY_VALIDATION_FAILURES_EXPRESSION
), **session._current_send_kwargs)
failure_message = e.message
if failure_message is None:
failure_message = render_expression(
config.DEFAULT_VALIDATION_FAILURE_EXPRESSION
)
config.DEFAULT_VALIDATION_FAILURE_EXPRESSION)
# noinspection PyProtectedMember
session.pause(failure_message,
**session._current_send_kwargs)
@ -157,7 +150,8 @@ class CommandFunc:
return parser_func
def on_command(name: Union[str, CommandName_T], *,
def on_command(name: Union[str, CommandName_T],
*,
aliases: Union[Iterable[str], str] = (),
permission: int = perm.EVERYBODY,
only_to_me: bool = True,
@ -173,18 +167,21 @@ def on_command(name: Union[str, CommandName_T], *,
:param privileged: can be run even when there is already a session
:param shell_like: use shell-like syntax to split arguments
"""
def deco(func: CommandHandler_T) -> CommandHandler_T:
if not isinstance(name, (str, tuple)):
raise TypeError('the name of a command must be a str or tuple')
if not name:
raise ValueError('the name of a command must not be empty')
cmd_name = (name,) if isinstance(name, str) else name
cmd_name = (name, ) if isinstance(name, str) else name
cmd = Command(name=cmd_name, func=func, permission=permission,
only_to_me=only_to_me, privileged=privileged)
cmd = Command(name=cmd_name,
func=func,
permission=permission,
only_to_me=only_to_me,
privileged=privileged)
if shell_like:
async def shell_like_args_parser(session):
session.args['argv'] = shlex.split(session.current_arg)
@ -204,7 +201,7 @@ def on_command(name: Union[str, CommandName_T], *,
nonlocal aliases
if isinstance(aliases, str):
aliases = (aliases,)
aliases = (aliases, )
for alias in aliases:
_aliases[alias] = cmd_name
@ -214,7 +211,7 @@ def on_command(name: Union[str, CommandName_T], *,
def _find_command(name: Union[str, CommandName_T]) -> Optional[Command]:
cmd_name = (name,) if isinstance(name, str) else name
cmd_name = (name, ) if isinstance(name, str) else name
if not cmd_name:
return None
@ -241,7 +238,6 @@ class _FinishException(Exception):
Raised by session.finish() indicating that the command session
should be stopped and removed.
"""
def __init__(self, result: bool = True):
"""
:param result: succeeded to call the command
@ -259,7 +255,6 @@ class SwitchException(Exception):
again, the later function should be notified. So this exception
is designed to be propagated to handle_message().
"""
def __init__(self, new_ctx_message: Message):
"""
:param new_ctx_message: new message which should be placed in context
@ -268,13 +263,18 @@ class SwitchException(Exception):
class CommandSession(BaseSession):
__slots__ = ('cmd',
'current_key', 'current_arg_filters', '_current_send_kwargs',
'current_arg', '_current_arg_text', '_current_arg_images',
'_state', '_last_interaction', '_running')
__slots__ = ('cmd', 'current_key', 'current_arg_filters',
'_current_send_kwargs', 'current_arg', '_current_arg_text',
'_current_arg_images', '_state', '_last_interaction',
'_running', '_run_future')
def __init__(self, bot: NoneBot, ctx: Context_T, cmd: Command, *,
current_arg: str = '', args: Optional[CommandArgs_T] = None):
def __init__(self,
bot: NoneBot,
ctx: Context_T,
cmd: Command,
*,
current_arg: str = '',
args: Optional[CommandArgs_T] = None):
super().__init__(bot, ctx)
self.cmd = cmd # Command object
@ -296,6 +296,9 @@ class CommandSession(BaseSession):
if args:
self._state.update(args)
self._run_future = partial(asyncio.run_coroutine_threadsafe,
loop=bot.loop)
self._last_interaction = None # last interaction time of this session
self._running = False
@ -381,7 +384,9 @@ class CommandSession(BaseSession):
self._current_arg_text = None
self._current_arg_images = None
def get(self, key: str, *,
def get(self,
key: str,
*,
prompt: Optional[Message_T] = None,
arg_filters: Optional[List[Filter_T]] = None,
**kwargs) -> Any:
@ -406,7 +411,8 @@ class CommandSession(BaseSession):
self._current_send_kwargs = kwargs
self.pause(prompt, **kwargs)
def get_optional(self, key: str,
def get_optional(self,
key: str,
default: Optional[Any] = None) -> Optional[Any]:
"""
Simply get a argument with given key.
@ -418,13 +424,13 @@ class CommandSession(BaseSession):
def pause(self, message: Optional[Message_T] = None, **kwargs) -> None:
"""Pause the session for further interaction."""
if message:
asyncio.ensure_future(self.send(message, **kwargs))
self._run_future(self.send(message, **kwargs))
raise _PauseException
def finish(self, message: Optional[Message_T] = None, **kwargs) -> None:
"""Finish the session."""
if message:
asyncio.ensure_future(self.send(message, **kwargs))
self._run_future(self.send(message, **kwargs))
raise _FinishException
def switch(self, new_ctx_message: Message_T) -> None:
@ -507,7 +513,7 @@ def parse_command(bot: NoneBot,
cmd_name = curr_cmd_name
if not cmd_name:
cmd_name = (cmd_name_text,)
cmd_name = (cmd_name_text, )
logger.debug(f'Split command name: {cmd_name}')
cmd = _find_command(cmd_name)
@ -554,10 +560,9 @@ async def handle_command(bot: NoneBot, ctx: Context_T) -> bool:
if session.running:
logger.warning(f'There is a session of command '
f'{session.cmd.name} running, notify the user')
asyncio.ensure_future(send(
bot, ctx,
render_expression(bot.config.SESSION_RUNNING_EXPRESSION)
))
asyncio.ensure_future(
send(bot, ctx,
render_expression(bot.config.SESSION_RUNNING_EXPRESSION)))
# pretend we are successful, so that NLP won't handle it
return True
@ -585,12 +590,16 @@ async def handle_command(bot: NoneBot, ctx: Context_T) -> bool:
session = CommandSession(bot, ctx, cmd, current_arg=current_arg)
logger.debug(f'New session of command {session.cmd.name} created')
return await _real_run_command(session, ctx_id, check_perm=check_perm,
return await _real_run_command(session,
ctx_id,
check_perm=check_perm,
disable_interaction=disable_interaction)
async def call_command(bot: NoneBot, ctx: Context_T,
name: Union[str, CommandName_T], *,
async def call_command(bot: NoneBot,
ctx: Context_T,
name: Union[str, CommandName_T],
*,
current_arg: str = '',
args: Optional[CommandArgs_T] = None,
check_perm: bool = True,
@ -619,7 +628,8 @@ async def call_command(bot: NoneBot, ctx: Context_T,
if not cmd:
return False
session = CommandSession(bot, ctx, cmd, current_arg=current_arg, args=args)
return await _real_run_command(session, context_id(session.ctx),
return await _real_run_command(session,
context_id(session.ctx),
check_perm=check_perm,
disable_interaction=disable_interaction)