#!/usr/bin/env python3 # -*- coding: utf-8 -*- import typing import inspect from functools import wraps from datetime import datetime from collections import defaultdict from contextvars import Context, ContextVar, copy_context 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 from nonebot.exception import PausedException, RejectedException, FinishedException matchers: Dict[int, List[Type["Matcher"]]] = defaultdict(list) current_bot: ContextVar = ContextVar("current_bot") current_event: ContextVar = ContextVar("current_event") class Matcher: """`Matcher`类 """ type: str = "" rule: Rule = Rule() permission: Permission = Permission() handlers: List[Handler] = [] temp: bool = False expire_time: Optional[datetime] = None priority: int = 1 block: bool = False _default_state: dict = {} _default_parser: Optional[ArgsParser] = None def __init__(self): """实例化 Matcher 以便运行 """ self.handlers = self.handlers.copy() self.state = self._default_state.copy() def __repr__(self) -> str: return (f"") def __str__(self) -> str: return self.__repr__() @classmethod def new(cls, type_: str = "", rule: Rule = Rule(), permission: Permission = Permission(), handlers: Optional[list] = None, temp: bool = False, priority: int = 1, block: bool = False, *, default_state: Optional[dict] = None, expire_time: Optional[datetime] = None) -> Type["Matcher"]: """创建新的 Matcher Returns: Type["Matcher"]: 新的 Matcher 类 """ NewMatcher = type( "Matcher", (Matcher,), { "type": type_, "rule": rule, "permission": permission, "handlers": handlers or [], "temp": temp, "expire_time": expire_time, "priority": priority, "block": block, "_default_state": default_state or {} }) matchers[priority].append(NewMatcher) return NewMatcher @classmethod async def check_perm(cls, bot: Bot, event: Event) -> bool: return (event.type == (cls.type or event.type) and await cls.permission(bot, event)) @classmethod async def check_rule(cls, bot: Bot, event: Event, state: dict) -> bool: """检查 Matcher 的 Rule 是否成立 Args: event (Event): 消息事件 Returns: bool: 条件成立与否 """ return await cls.rule(bot, event, state) @classmethod def args_parser(cls, func: ArgsParser) -> ArgsParser: cls._default_parser = func return func @classmethod def handle(cls) -> Callable[[Handler], Handler]: """直接处理消息事件""" def _decorator(func: Handler) -> Handler: cls.handlers.append(func) return func return _decorator @classmethod def receive(cls) -> Callable[[Handler], Handler]: """接收一条新消息并处理""" async def _receive(bot: Bot, event: Event, state: dict) -> NoReturn: raise PausedException if cls.handlers: # 已有前置handlers则接受一条新的消息,否则视为接收初始消息 cls.handlers.append(_receive) def _decorator(func: Handler) -> Handler: if not cls.handlers or cls.handlers[-1] is not func: cls.handlers.append(func) return func return _decorator @classmethod def got( cls, key: str, prompt: Optional[str] = None, args_parser: Optional[ArgsParser] = None ) -> Callable[[Handler], Handler]: 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) raise PausedException async def _key_parser(bot: Bot, event: Event, state: dict): if key in state: return parser = args_parser or cls._default_parser if parser: await parser(bot, event, state) else: state[state["_current_key"]] = str(event.message) cls.handlers.append(_key_getter) cls.handlers.append(_key_parser) def _decorator(func: Handler) -> Handler: 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) return func return _decorator @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 @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 @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 # 运行handlers async def run(self, bot: Bot, event: Event, state: dict): b_t = current_bot.set(bot) e_t = current_event.set(event) try: # Refresh preprocess state self.state.update(state) for _ in range(len(self.handlers)): handler = self.handlers.pop(0) annotation = typing.get_type_hints(handler) BotType = annotation.get("bot") if BotType and inspect.isclass(BotType) and not isinstance( bot, BotType): continue await handler(bot, event, self.state) except RejectedException: self.handlers.insert(0, handler) # type: ignore Matcher.new( self.type, Rule(), USER(event.user_id, perm=self.permission), # type:ignore self.handlers, temp=True, priority=0, block=True, default_state=self.state, expire_time=datetime.now() + bot.config.session_expire_timeout) except PausedException: Matcher.new( self.type, Rule(), USER(event.user_id, perm=self.permission), # type:ignore self.handlers, temp=True, priority=0, block=True, default_state=self.state, expire_time=datetime.now() + bot.config.session_expire_timeout) except FinishedException: pass finally: current_bot.reset(b_t) current_event.reset(e_t)