nonebot2/nonebot/matcher.py

274 lines
8.4 KiB
Python
Raw Normal View History

2020-06-30 10:13:58 +08:00
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
2020-08-27 16:43:58 +08:00
from nonebot.log import logger
2020-08-13 15:23:04 +08:00
import typing
2020-08-24 17:59:36 +08:00
import inspect
2020-05-02 20:03:36 +08:00
from functools import wraps
2020-08-06 17:54:55 +08:00
from datetime import datetime
2020-06-30 10:13:58 +08:00
from collections import defaultdict
from contextvars import Context, ContextVar, copy_context
2020-05-02 20:03:36 +08:00
2020-08-17 16:09:41 +08:00
from nonebot.rule import Rule
from nonebot.permission import Permission, USER
from nonebot.typing import Type, List, Dict, Union, Callable, Optional, NoReturn
from nonebot.typing import Bot, Event, Handler, Message, ArgsParser, MessageSegment
2020-07-05 20:39:34 +08:00
from nonebot.exception import PausedException, RejectedException, FinishedException
2020-05-02 20:03:36 +08:00
2020-06-30 10:13:58 +08:00
matchers: Dict[int, List[Type["Matcher"]]] = defaultdict(list)
current_bot: ContextVar = ContextVar("current_bot")
current_event: ContextVar = ContextVar("current_event")
2020-06-30 10:13:58 +08:00
2020-05-02 20:03:36 +08:00
2020-08-27 16:43:58 +08:00
class MatcherMeta(type):
def __repr__(self) -> str:
return (f"<Matcher from {self.module or 'unknow'}, " # type: ignore
f"type={self.type}, priority={self.priority}, " # type: ignore
f"temp={self.temp}>") # type: ignore
def __str__(self) -> str:
2020-08-28 11:54:21 +08:00
return repr(self)
2020-08-27 16:43:58 +08:00
class Matcher(metaclass=MatcherMeta):
2020-07-04 22:51:10 +08:00
"""`Matcher`类
"""
2020-08-27 16:43:58 +08:00
module: Optional[str] = None
2020-05-02 20:03:36 +08:00
2020-08-20 17:15:05 +08:00
type: str = ""
2020-08-17 16:09:41 +08:00
rule: Rule = Rule()
permission: Permission = Permission()
2020-06-30 10:13:58 +08:00
handlers: List[Handler] = []
2020-05-05 16:11:05 +08:00
temp: bool = False
2020-08-06 17:54:55 +08:00
expire_time: Optional[datetime] = None
2020-06-30 10:13:58 +08:00
priority: int = 1
block: bool = False
2020-05-05 16:11:05 +08:00
_default_state: dict = {}
2020-06-30 10:13:58 +08:00
_default_parser: Optional[ArgsParser] = None
2020-05-05 16:11:05 +08:00
def __init__(self):
2020-07-04 22:51:10 +08:00
"""实例化 Matcher 以便运行
"""
2020-05-05 16:11:05 +08:00
self.handlers = self.handlers.copy()
self.state = self._default_state.copy()
2020-08-24 17:59:36 +08:00
def __repr__(self) -> str:
2020-08-27 16:43:58 +08:00
return (f"<Matcher from {self.module or 'unknow'}, type={self.type}, "
f"priority={self.priority}, temp={self.temp}>")
2020-08-24 17:59:36 +08:00
def __str__(self) -> str:
return self.__repr__()
2020-05-05 16:11:05 +08:00
@classmethod
2020-06-30 10:13:58 +08:00
def new(cls,
2020-08-20 17:15:05 +08:00
type_: str = "",
2020-08-17 16:09:41 +08:00
rule: Rule = Rule(),
permission: Permission = Permission(),
handlers: Optional[list] = None,
2020-06-30 10:13:58 +08:00
temp: bool = False,
priority: int = 1,
block: bool = False,
2020-06-30 10:13:58 +08:00
*,
2020-08-27 16:43:58 +08:00
module: Optional[str] = None,
default_state: Optional[dict] = None,
2020-08-06 17:54:55 +08:00
expire_time: Optional[datetime] = None) -> Type["Matcher"]:
2020-07-04 22:51:10 +08:00
"""创建新的 Matcher
Returns:
Type["Matcher"]: 新的 Matcher
"""
2020-05-05 16:11:05 +08:00
NewMatcher = type(
2020-06-24 21:56:49 +08:00
"Matcher", (Matcher,), {
2020-08-27 16:43:58 +08:00
"module": module,
2020-08-20 17:15:05 +08:00
"type": type_,
2020-05-05 16:11:05 +08:00
"rule": rule,
2020-08-17 16:09:41 +08:00
"permission": permission,
"handlers": handlers or [],
2020-05-05 16:11:05 +08:00
"temp": temp,
2020-08-06 17:54:55 +08:00
"expire_time": expire_time,
2020-06-30 10:13:58 +08:00
"priority": priority,
"block": block,
"_default_state": default_state or {}
2020-05-05 16:11:05 +08:00
})
2020-06-30 10:13:58 +08:00
matchers[priority].append(NewMatcher)
2020-05-05 16:11:05 +08:00
return NewMatcher
2020-07-04 22:51:10 +08:00
@classmethod
2020-08-17 16:09:41 +08:00
async def check_perm(cls, bot: Bot, event: Event) -> bool:
2020-08-20 17:15:05 +08:00
return (event.type == (cls.type or event.type) and
await cls.permission(bot, event))
2020-08-17 16:09:41 +08:00
@classmethod
async def check_rule(cls, bot: Bot, event: Event, state: dict) -> bool:
2020-07-04 22:51:10 +08:00
"""检查 Matcher 的 Rule 是否成立
Args:
event (Event): 消息事件
Returns:
bool: 条件成立与否
"""
2020-08-17 16:09:41 +08:00
return await cls.rule(bot, event, state)
2020-07-04 22:51:10 +08:00
@classmethod
def args_parser(cls, func: ArgsParser) -> ArgsParser:
cls._default_parser = func
return func
2020-06-30 10:13:58 +08:00
2020-05-05 16:11:05 +08:00
@classmethod
def handle(cls) -> Callable[[Handler], Handler]:
2020-07-04 22:51:10 +08:00
"""直接处理消息事件"""
2020-06-30 10:13:58 +08:00
def _decorator(func: Handler) -> Handler:
cls.handlers.append(func)
return func
return _decorator
2020-05-02 20:03:36 +08:00
2020-05-05 16:11:05 +08:00
@classmethod
def receive(cls) -> Callable[[Handler], Handler]:
2020-07-04 22:51:10 +08:00
"""接收一条新消息并处理"""
2020-05-02 20:03:36 +08:00
async def _receive(bot: Bot, event: Event, state: dict) -> NoReturn:
2020-08-22 15:22:46 +08:00
raise PausedException
2020-05-02 20:03:36 +08:00
2020-08-22 15:22:46 +08:00
if cls.handlers:
# 已有前置handlers则接受一条新的消息否则视为接收初始消息
cls.handlers.append(_receive)
2020-05-05 16:11:05 +08:00
2020-08-22 15:22:46 +08:00
def _decorator(func: Handler) -> Handler:
2020-08-24 17:59:36 +08:00
if not cls.handlers or cls.handlers[-1] is not func:
2020-08-22 15:22:46 +08:00
cls.handlers.append(func)
2020-05-02 20:03:36 +08:00
return func
return _decorator
@classmethod
def got(
cls,
key: str,
prompt: Optional[str] = None,
args_parser: Optional[ArgsParser] = None
) -> Callable[[Handler], Handler]:
2020-05-02 20:03:36 +08:00
2020-08-22 15:22:46 +08:00
async def _key_getter(bot: Bot, event: Event, state: dict):
if key not in state:
state["_current_key"] = key
if prompt:
await bot.send(event=event, message=prompt)
2020-08-22 15:22:46 +08:00
raise PausedException
2020-05-02 20:03:36 +08:00
2020-08-22 15:22:46 +08:00
async def _key_parser(bot: Bot, event: Event, state: dict):
2020-08-24 17:59:36 +08:00
if key in state:
return
2020-08-22 15:22:46 +08:00
parser = args_parser or cls._default_parser
if parser:
await parser(bot, event, state)
else:
state[state["_current_key"]] = str(event.message)
2020-08-23 10:45:26 +08:00
cls.handlers.append(_key_getter)
2020-08-22 15:22:46 +08:00
cls.handlers.append(_key_parser)
def _decorator(func: Handler) -> Handler:
2020-08-24 17:59:36 +08:00
if not hasattr(cls.handlers[-1], "__wrapped__"):
parser = cls.handlers.pop()
@wraps(func)
async def wrapper(bot: Bot, event: Event, state: dict):
await parser(bot, event, state)
await func(bot, event, state)
cls.handlers.append(wrapper)
2020-05-02 20:03:36 +08:00
return func
2020-05-02 20:03:36 +08:00
return _decorator
2020-05-02 20:03:36 +08:00
@classmethod
async def finish(
cls,
prompt: Optional[Union[str, Message,
MessageSegment]] = None) -> NoReturn:
bot: Bot = current_bot.get()
event: Event = current_event.get()
if prompt:
await bot.send(event=event, message=prompt)
raise FinishedException
2020-05-02 20:03:36 +08:00
@classmethod
async def pause(
cls,
prompt: Optional[Union[str, Message,
MessageSegment]] = None) -> NoReturn:
bot: Bot = current_bot.get()
event: Event = current_event.get()
if prompt:
await bot.send(event=event, message=prompt)
raise PausedException
2020-06-30 10:13:58 +08:00
@classmethod
async def reject(
cls,
prompt: Optional[Union[str, Message,
MessageSegment]] = None) -> NoReturn:
bot: Bot = current_bot.get()
event: Event = current_event.get()
if prompt:
await bot.send(event=event, message=prompt)
raise RejectedException
2020-05-02 20:03:36 +08:00
2020-07-04 22:51:10 +08:00
# 运行handlers
async def run(self, bot: Bot, event: Event, state: dict):
b_t = current_bot.set(bot)
e_t = current_event.set(event)
2020-05-05 16:11:05 +08:00
try:
2020-08-17 16:09:41 +08:00
# Refresh preprocess state
self.state.update(state)
2020-05-05 16:11:05 +08:00
for _ in range(len(self.handlers)):
handler = self.handlers.pop(0)
2020-08-13 15:23:04 +08:00
annotation = typing.get_type_hints(handler)
BotType = annotation.get("bot")
2020-08-24 17:59:36 +08:00
if BotType and inspect.isclass(BotType) and not isinstance(
bot, BotType):
2020-08-13 15:23:04 +08:00
continue
2020-06-30 10:13:58 +08:00
await handler(bot, event, self.state)
2020-05-05 16:11:05 +08:00
except RejectedException:
2020-06-30 10:13:58 +08:00
self.handlers.insert(0, handler) # type: ignore
2020-08-24 17:59:36 +08:00
Matcher.new(
2020-08-20 17:15:05 +08:00
self.type,
2020-08-24 17:59:36 +08:00
Rule(),
2020-08-17 16:09:41 +08:00
USER(event.user_id, perm=self.permission), # type:ignore
self.handlers,
temp=True,
priority=0,
2020-08-24 17:59:36 +08:00
block=True,
2020-08-17 16:09:41 +08:00
default_state=self.state,
expire_time=datetime.now() + bot.config.session_expire_timeout)
2020-05-05 16:11:05 +08:00
except PausedException:
2020-08-24 17:59:36 +08:00
Matcher.new(
2020-08-20 17:15:05 +08:00
self.type,
2020-08-24 17:59:36 +08:00
Rule(),
2020-08-17 16:09:41 +08:00
USER(event.user_id, perm=self.permission), # type:ignore
self.handlers,
temp=True,
priority=0,
2020-08-24 17:59:36 +08:00
block=True,
2020-08-17 16:09:41 +08:00
default_state=self.state,
expire_time=datetime.now() + bot.config.session_expire_timeout)
2020-05-05 16:11:05 +08:00
except FinishedException:
pass
finally:
2020-08-27 16:43:58 +08:00
logger.info(f"Matcher {self} running complete")
current_bot.reset(b_t)
current_event.reset(e_t)