mirror of
https://github.com/nonebot/nonebot2.git
synced 2024-11-24 17:15:05 +08:00
278 lines
9.1 KiB
Python
278 lines
9.1 KiB
Python
import re
|
||
from collections import defaultdict
|
||
from typing import Tuple, Union, Callable, Iterable, Dict, Any, Optional
|
||
|
||
from aiocqhttp import CQHttp, Error as CQHttpError
|
||
from aiocqhttp.message import Message
|
||
|
||
from . import permissions as perm
|
||
from .helpers import context_source
|
||
|
||
# Key: str (one segment of command name)
|
||
# Value: subtree or a leaf Command object
|
||
_registry = {}
|
||
|
||
# Key: str
|
||
# Value: tuple that identifies a command
|
||
_aliases = {}
|
||
|
||
# Key: context source
|
||
# Value: list (stack) of Session objects
|
||
_sessions = defaultdict(list)
|
||
|
||
|
||
class Command:
|
||
__slots__ = ('name', 'func', 'permission', 'args_parser')
|
||
|
||
def __init__(self, name: Tuple[str], func: Callable, permission: int):
|
||
self.name = name
|
||
self.func = func
|
||
self.permission = permission
|
||
self.args_parser = None
|
||
|
||
async def run(self, bot, session, *,
|
||
permission: int = None) -> bool:
|
||
if permission is None:
|
||
permission = await calculate_permission(bot, session.ctx)
|
||
if isinstance(self.func, Callable) and permission & self.permission:
|
||
if isinstance(self.args_parser, Callable):
|
||
self.args_parser(session)
|
||
await self.func(bot, session)
|
||
return True
|
||
return False
|
||
|
||
|
||
async def calculate_permission(bot: CQHttp, ctx: Dict[str, Any]) -> int:
|
||
permission = 0
|
||
if ctx['user_id'] in bot.config.SUPERUSERS:
|
||
permission |= perm.IS_SUPERUSER
|
||
if ctx['message_type'] == 'private':
|
||
if ctx['sub_type'] == 'friend':
|
||
permission |= perm.IS_PRIVATE_FRIEND
|
||
elif ctx['sub_type'] == 'group':
|
||
permission |= perm.IS_PRIVATE_GROUP
|
||
elif ctx['sub_type'] == 'discuss':
|
||
permission |= perm.IS_PRIVATE_DISCUSS
|
||
elif ctx['sub_type'] == 'other':
|
||
permission |= perm.IS_PRIVATE_OTHER
|
||
elif ctx['message_type'] == 'group':
|
||
permission |= perm.IS_GROUP_MEMBER
|
||
if not ctx['anonymous']:
|
||
try:
|
||
member_info = await bot.get_group_member_info(**ctx)
|
||
if member_info:
|
||
if member_info['role'] == 'owner':
|
||
permission |= perm.IS_GROUP_OWNER
|
||
elif member_info['role'] == 'admin':
|
||
permission |= perm.IS_GROUP_ADMIN
|
||
except CQHttpError:
|
||
pass
|
||
elif ctx['message_type'] == 'discuss':
|
||
permission |= perm.IS_DISCUSS
|
||
return permission
|
||
|
||
|
||
def on_command(name: Union[str, Tuple[str]], aliases: Iterable = (),
|
||
permission: int = perm.EVERYONE) -> Callable:
|
||
def deco(func: Callable) -> Callable:
|
||
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, tuple) else (name,)
|
||
current_parent = _registry
|
||
for parent_key in cmd_name[:-1]:
|
||
current_parent[parent_key] = {}
|
||
current_parent = current_parent[parent_key]
|
||
cmd = Command(name=cmd_name, func=func, permission=permission)
|
||
current_parent[cmd_name[-1]] = cmd
|
||
for alias in aliases:
|
||
_aliases[alias] = cmd_name
|
||
|
||
def args_parser(parser_func: Callable):
|
||
cmd.args_parser = parser_func
|
||
return parser_func
|
||
|
||
func.args_parser = args_parser
|
||
return func
|
||
|
||
return deco
|
||
|
||
|
||
def _find_command(name: Union[str, Tuple[str]]) -> Optional[Command]:
|
||
cmd_name = name if isinstance(name, tuple) else (name,)
|
||
|
||
if not cmd_name:
|
||
return None
|
||
|
||
cmd_tree = _registry
|
||
for part in cmd_name[:-1]:
|
||
if part not in cmd_tree:
|
||
return None
|
||
cmd_tree = cmd_tree[part]
|
||
|
||
return cmd_tree.get(cmd_name[-1])
|
||
|
||
|
||
class FurtherInteractionNeeded(Exception):
|
||
"""
|
||
Raised by session.require_arg() indicating
|
||
that the command should enter interactive mode
|
||
to ask the user for some arguments.
|
||
"""
|
||
pass
|
||
|
||
|
||
class Session:
|
||
__slots__ = ('cmd', 'ctx',
|
||
'current_key', 'current_prompt',
|
||
'current_arg', 'current_arg_text',
|
||
'images', 'args', 'last_interaction')
|
||
|
||
def __init__(self, cmd: Command, ctx: Dict[str, Any], *,
|
||
current_arg: str = '', args: Dict[str, Any] = None):
|
||
self.cmd = cmd
|
||
self.ctx = ctx
|
||
self.current_key = None
|
||
self.current_prompt = None
|
||
self.current_arg = current_arg
|
||
self.current_arg_text = Message(current_arg).extract_plain_text()
|
||
self.images = [s.data['url'] for s in ctx['message']
|
||
if s.type == 'image' and 'url' in s.data]
|
||
self.args = args or {}
|
||
self.last_interaction = None
|
||
|
||
def refresh(self, ctx: Dict[str, Any], *, current_arg: str = ''):
|
||
self.ctx = ctx
|
||
self.current_arg = current_arg
|
||
self.current_arg_text = Message(current_arg).extract_plain_text()
|
||
self.images = [s.data['url'] for s in ctx['message']
|
||
if s.type == 'image' and 'url' in s.data]
|
||
|
||
@property
|
||
def is_valid(self):
|
||
# TODO: 检查 last_interaction
|
||
return True
|
||
|
||
def require_arg(self, key: str, prompt: str = None, *,
|
||
interactive: bool = True) -> Any:
|
||
"""
|
||
Get an argument with a given key.
|
||
|
||
If "interactive" is True, and the argument does not exist
|
||
in the current session, a FurtherInteractionNeeded exception
|
||
will be raised, and the caller of the command will know
|
||
it should keep the session for further interaction with the user.
|
||
|
||
If "interactive" is False, missed key will cause a result of None.
|
||
|
||
:param key: argument key
|
||
:param prompt: prompt to ask the user with
|
||
:param interactive: should enter interactive mode while key missing
|
||
:return: the argument value
|
||
:raise FurtherInteractionNeeded: further interaction is needed
|
||
"""
|
||
value = self.args.get(key)
|
||
if value is not None or not interactive:
|
||
return value
|
||
|
||
self.current_key = key
|
||
self.current_prompt = prompt or f'请输入 {self.current_key}：'
|
||
raise FurtherInteractionNeeded
|
||
|
||
|
||
def _new_command_session(bot: CQHttp,
|
||
ctx: Dict[str, Any]) -> Optional[Session]:
|
||
msg_text = str(ctx['message']).lstrip()
|
||
|
||
for start in bot.config.COMMAND_START:
|
||
if isinstance(start, type(re.compile(''))):
|
||
m = start.search(msg_text)
|
||
if m:
|
||
full_command = msg_text[len(m.group(0)):].lstrip()
|
||
break
|
||
elif isinstance(start, str):
|
||
if msg_text.startswith(start):
|
||
full_command = msg_text[len(start):].lstrip()
|
||
break
|
||
else:
|
||
# it's not a command
|
||
return None
|
||
|
||
if not full_command:
|
||
# command is empty
|
||
return None
|
||
|
||
cmd_name_text, *cmd_remained = full_command.split(maxsplit=1)
|
||
cmd_name = _aliases.get(cmd_name_text)
|
||
|
||
if not cmd_name:
|
||
for sep in bot.config.COMMAND_SEP:
|
||
if isinstance(sep, type(re.compile(''))):
|
||
cmd_name = tuple(sep.split(cmd_name_text))
|
||
break
|
||
elif isinstance(sep, str):
|
||
cmd_name = tuple(cmd_name_text.split(sep))
|
||
break
|
||
else:
|
||
cmd_name = (cmd_name_text,)
|
||
|
||
cmd = _find_command(cmd_name)
|
||
if not cmd:
|
||
return None
|
||
|
||
return Session(cmd, ctx, current_arg=''.join(cmd_remained))
|
||
|
||
|
||
async def handle_command(bot: CQHttp, ctx: Dict[str, Any]) -> bool:
|
||
src = context_source(ctx)
|
||
if _sessions[src]:
|
||
session = _sessions[src][-1]
|
||
session.refresh(ctx, current_arg=str(ctx['message']))
|
||
# TODO: 检查 is_valid
|
||
else:
|
||
session = _new_command_session(bot, ctx)
|
||
if not session:
|
||
return False
|
||
_sessions[src].append(session)
|
||
|
||
try:
|
||
res = await session.cmd.run(bot, session)
|
||
# the command is finished, pop the session
|
||
_sessions[src].pop()
|
||
if not _sessions[src]:
|
||
# session stack of the current user is empty
|
||
del _sessions[src]
|
||
return res
|
||
except FurtherInteractionNeeded:
|
||
# ask the user for more information
|
||
await bot.send(ctx, session.current_prompt)
|
||
|
||
# return True because this step of the session is successful
|
||
return True
|
||
|
||
|
||
async def call_command(name: Union[str, Tuple[str]],
|
||
bot: CQHttp, ctx: Dict[str, Any], **kwargs) -> bool:
|
||
"""
|
||
Call a command internally.
|
||
|
||
There is no permission restriction on this function,
|
||
which means any command can be called from any other command.
|
||
Unexpected users should be handled by the caller command's permission
|
||
option.
|
||
|
||
:param name: command name (str or tuple of str)
|
||
:param bot: CQHttp instance
|
||
:param ctx: event context
|
||
:param kwargs: other keyword args that will be passed to Session()
|
||
:return: the command is successfully called
|
||
"""
|
||
cmd = _find_command(name)
|
||
if cmd:
|
||
session = Session(cmd, ctx, **kwargs)
|
||
# TODO: 插入 session
|
||
return await cmd.run(bot, session, permission=perm.IS_SUPERUSER)
|
||
return False
|