mirror of
https://github.com/nonebot/nonebot2.git
synced 2025-07-16 19:11:00 +00:00
add permission and command
This commit is contained in:
yanyongyu
281
nonebot/rule.py
281
nonebot/rule.py
@ -2,207 +2,126 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
import re
|
||||
import abc
|
||||
import asyncio
|
||||
from typing import cast
|
||||
from itertools import product
|
||||
|
||||
from pygtrie import CharTrie
|
||||
|
||||
from nonebot import get_driver
|
||||
from nonebot.log import logger
|
||||
from nonebot.utils import run_sync
|
||||
from nonebot.typing import Bot, Event, Union, Optional, Awaitable
|
||||
from nonebot.typing import RuleChecker, SyncRuleChecker, AsyncRuleChecker
|
||||
from nonebot.typing import Bot, Any, Dict, Event, Union, Tuple, NoReturn, RuleChecker
|
||||
|
||||
|
||||
class BaseRule(abc.ABC):
|
||||
class Rule:
|
||||
__slots__ = ("checkers",)
|
||||
|
||||
def __init__(self, checker: RuleChecker):
|
||||
self.checker: RuleChecker = checker
|
||||
def __init__(self, *checkers: RuleChecker) -> None:
|
||||
self.checkers = list(checkers)
|
||||
|
||||
@abc.abstractmethod
|
||||
def __call__(self, bot: Bot, event: Event) -> Awaitable[bool]:
|
||||
raise NotImplementedError
|
||||
async def __call__(self, bot: Bot, event: Event, state: dict) -> bool:
|
||||
results = await asyncio.gather(
|
||||
*map(lambda c: c(bot, event, state), self.checkers))
|
||||
return all(results)
|
||||
|
||||
@abc.abstractmethod
|
||||
def __and__(self, other: Union["BaseRule", RuleChecker]) -> "BaseRule":
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def __or__(self, other: Union["BaseRule", RuleChecker]) -> "BaseRule":
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def __neg__(self) -> "BaseRule":
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AsyncRule(BaseRule):
|
||||
|
||||
def __init__(self, checker: Optional[AsyncRuleChecker] = None):
|
||||
|
||||
async def always_true(bot: Bot, event: Event) -> bool:
|
||||
return True
|
||||
|
||||
self.checker: AsyncRuleChecker = checker or always_true
|
||||
|
||||
def __call__(self, bot: Bot, event: Event) -> Awaitable[bool]:
|
||||
return self.checker(bot, event)
|
||||
|
||||
def __and__(self, other: Union[BaseRule, RuleChecker]) -> "AsyncRule":
|
||||
func = other
|
||||
if isinstance(other, BaseRule):
|
||||
func = other.checker
|
||||
|
||||
if not asyncio.iscoroutinefunction(func):
|
||||
func = run_sync(func)
|
||||
|
||||
async def tmp(bot: Bot, event: Event) -> bool:
|
||||
a, b = await asyncio.gather(self.checker(bot, event),
|
||||
func(bot, event))
|
||||
return a and b
|
||||
|
||||
return AsyncRule(tmp)
|
||||
|
||||
def __or__(self, other: Union[BaseRule, RuleChecker]) -> "AsyncRule":
|
||||
func = other
|
||||
if isinstance(other, BaseRule):
|
||||
func = other.checker
|
||||
|
||||
if not asyncio.iscoroutinefunction(func):
|
||||
func = run_sync(func)
|
||||
|
||||
async def tmp(bot: Bot, event: Event) -> bool:
|
||||
a, b = await asyncio.gather(self.checker(bot, event),
|
||||
func(bot, event))
|
||||
return a or b
|
||||
|
||||
return AsyncRule(tmp)
|
||||
|
||||
def __neg__(self) -> "AsyncRule":
|
||||
|
||||
async def neg(bot: Bot, event: Event) -> bool:
|
||||
result = await self.checker(bot, event)
|
||||
return not result
|
||||
|
||||
return AsyncRule(neg)
|
||||
|
||||
|
||||
class SyncRule(BaseRule):
|
||||
|
||||
def __init__(self, checker: Optional[SyncRuleChecker] = None):
|
||||
|
||||
def always_true(bot: Bot, event: Event) -> bool:
|
||||
return True
|
||||
|
||||
self.checker: SyncRuleChecker = checker or always_true
|
||||
|
||||
def __call__(self, bot: Bot, event: Event) -> Awaitable[bool]:
|
||||
return run_sync(self.checker)(bot, event)
|
||||
|
||||
def __and__(self, other: Union[BaseRule, RuleChecker]) -> BaseRule:
|
||||
func = other
|
||||
if isinstance(other, BaseRule):
|
||||
func = other.checker
|
||||
|
||||
if not asyncio.iscoroutinefunction(func):
|
||||
# func: SyncRuleChecker
|
||||
syncfunc = cast(SyncRuleChecker, func)
|
||||
|
||||
def tmp(bot: Bot, event: Event) -> bool:
|
||||
return self.checker(bot, event) and syncfunc(bot, event)
|
||||
|
||||
return SyncRule(tmp)
|
||||
def __and__(self, other: Union["Rule", RuleChecker]) -> "Rule":
|
||||
checkers = [*self.checkers]
|
||||
if isinstance(other, Rule):
|
||||
checkers.extend(other.checkers)
|
||||
elif asyncio.iscoroutinefunction(other):
|
||||
checkers.append(other)
|
||||
else:
|
||||
# func: AsyncRuleChecker
|
||||
asyncfunc = cast(AsyncRuleChecker, func)
|
||||
checkers.append(run_sync(other))
|
||||
return Rule(*checkers)
|
||||
|
||||
async def tmp(bot: Bot, event: Event) -> bool:
|
||||
a, b = await asyncio.gather(
|
||||
run_sync(self.checker)(bot, event), asyncfunc(bot, event))
|
||||
return a and b
|
||||
|
||||
return AsyncRule(tmp)
|
||||
|
||||
def __or__(self, other: Union[BaseRule, RuleChecker]) -> BaseRule:
|
||||
func = other
|
||||
if isinstance(other, BaseRule):
|
||||
func = other.checker
|
||||
|
||||
if not asyncio.iscoroutinefunction(func):
|
||||
# func: SyncRuleChecker
|
||||
syncfunc = cast(SyncRuleChecker, func)
|
||||
|
||||
def tmp(bot: Bot, event: Event) -> bool:
|
||||
return self.checker(bot, event) or syncfunc(bot, event)
|
||||
|
||||
return SyncRule(tmp)
|
||||
else:
|
||||
# func: AsyncRuleChecker
|
||||
asyncfunc = cast(AsyncRuleChecker, func)
|
||||
|
||||
async def tmp(bot: Bot, event: Event) -> bool:
|
||||
a, b = await asyncio.gather(
|
||||
run_sync(self.checker)(bot, event), asyncfunc(bot, event))
|
||||
return a or b
|
||||
|
||||
return AsyncRule(tmp)
|
||||
|
||||
def __neg__(self) -> "SyncRule":
|
||||
|
||||
def neg(bot: Bot, event: Event) -> bool:
|
||||
return not self.checker(bot, event)
|
||||
|
||||
return SyncRule(neg)
|
||||
def __or__(self, other) -> NoReturn:
|
||||
raise RuntimeError("Or operation between rules is not allowed.")
|
||||
|
||||
|
||||
def Rule(func: Optional[RuleChecker] = None) -> BaseRule:
|
||||
if func and asyncio.iscoroutinefunction(func):
|
||||
asyncfunc = cast(AsyncRuleChecker, func)
|
||||
return AsyncRule(asyncfunc)
|
||||
else:
|
||||
syncfunc = cast(Optional[SyncRuleChecker], func)
|
||||
return SyncRule(syncfunc)
|
||||
class TrieRule:
|
||||
prefix: CharTrie = CharTrie()
|
||||
suffix: CharTrie = CharTrie()
|
||||
|
||||
@classmethod
|
||||
def add_prefix(cls, prefix: str, value: Any):
|
||||
if prefix in cls.prefix:
|
||||
logger.warning(f'Duplicated prefix rule "{prefix}"')
|
||||
return
|
||||
cls.prefix[prefix] = value
|
||||
|
||||
@classmethod
|
||||
def add_suffix(cls, suffix: str, value: Any):
|
||||
if suffix[::-1] in cls.suffix:
|
||||
logger.warning(f'Duplicated suffix rule "{suffix}"')
|
||||
return
|
||||
cls.suffix[suffix[::-1]] = value
|
||||
|
||||
@classmethod
|
||||
def get_value(cls, bot: Bot, event: Event,
|
||||
state: dict) -> Tuple[Dict[str, Any], Dict[str, Any]]:
|
||||
prefix = None
|
||||
suffix = None
|
||||
message = event.message[0]
|
||||
if message.type == "text":
|
||||
prefix = cls.prefix.longest_prefix(message.data["text"].lstrip())
|
||||
message_r = event.message[-1]
|
||||
if message_r.type == "text":
|
||||
suffix = cls.suffix.longest_prefix(
|
||||
message_r.data["text"].rstrip()[::-1])
|
||||
|
||||
state["_prefix"] = {prefix.key: prefix.value} if prefix else {}
|
||||
state["_suffix"] = {suffix.key: suffix.value} if suffix else {}
|
||||
|
||||
return ({
|
||||
prefix.key: prefix.value
|
||||
} if prefix else {}, {
|
||||
suffix.key: suffix.value
|
||||
} if suffix else {})
|
||||
|
||||
|
||||
def message() -> BaseRule:
|
||||
return Rule(lambda bot, event: event.type == "message")
|
||||
def startswith(msg: str) -> Rule:
|
||||
TrieRule.add_prefix(msg, (msg,))
|
||||
|
||||
async def _startswith(bot: Bot, event: Event, state: dict) -> bool:
|
||||
return msg in state["_prefix"]
|
||||
|
||||
return Rule(_startswith)
|
||||
|
||||
|
||||
def notice() -> BaseRule:
|
||||
return Rule(lambda bot, event: event.type == "notice")
|
||||
def endswith(msg: str) -> Rule:
|
||||
TrieRule.add_suffix(msg, (msg,))
|
||||
|
||||
async def _endswith(bot: Bot, event: Event, state: dict) -> bool:
|
||||
return msg in state["_suffix"]
|
||||
|
||||
return Rule(_endswith)
|
||||
|
||||
|
||||
def request() -> BaseRule:
|
||||
return Rule(lambda bot, event: event.type == "request")
|
||||
def keyword(msg: str) -> Rule:
|
||||
|
||||
async def _keyword(bot: Bot, event: Event, state: dict) -> bool:
|
||||
return bool(event.plain_text and msg in event.plain_text)
|
||||
|
||||
return Rule(_keyword)
|
||||
|
||||
|
||||
def metaevent() -> BaseRule:
|
||||
return Rule(lambda bot, event: event.type == "meta_event")
|
||||
def command(command: Tuple[str]) -> Rule:
|
||||
config = get_driver().config
|
||||
command_start = config.command_start
|
||||
command_sep = config.command_sep
|
||||
for start, sep in product(command_start, command_sep):
|
||||
TrieRule.add_prefix(f"{start}{sep.join(command)}", command)
|
||||
|
||||
async def _command(bot: Bot, event: Event, state: dict) -> bool:
|
||||
return command in state["_prefix"].values()
|
||||
|
||||
return Rule(_command)
|
||||
|
||||
|
||||
def user(*qq: int) -> BaseRule:
|
||||
return Rule(lambda bot, event: event.user_id in qq)
|
||||
|
||||
|
||||
def private() -> BaseRule:
|
||||
return Rule(lambda bot, event: event.detail_type == "private")
|
||||
|
||||
|
||||
def group(*group: int) -> BaseRule:
|
||||
return Rule(lambda bot, event: event.detail_type == "group" and event.
|
||||
group_id in group)
|
||||
|
||||
|
||||
def startswith(msg, start: int = None, end: int = None) -> BaseRule:
|
||||
return Rule(lambda bot, event: event.message.startswith(msg, start, end))
|
||||
|
||||
|
||||
def endswith(msg, start: int = None, end: int = None) -> BaseRule:
|
||||
return Rule(
|
||||
lambda bot, event: event.message.endswith(msg, start=None, end=None))
|
||||
|
||||
|
||||
def has(msg: str) -> BaseRule:
|
||||
return Rule(lambda bot, event: msg in event.message)
|
||||
|
||||
|
||||
def regex(regex, flags: Union[int, re.RegexFlag] = 0) -> BaseRule:
|
||||
def regex(regex: str, flags: Union[int, re.RegexFlag] = 0) -> Rule:
|
||||
pattern = re.compile(regex, flags)
|
||||
return Rule(lambda bot, event: bool(pattern.search(str(event.message))))
|
||||
|
||||
async def _regex(bot: Bot, event: Event, state: dict) -> bool:
|
||||
return bool(pattern.search(str(event.message)))
|
||||
|
||||
return Rule(_regex)
|
||||
|
||||
Reference in New Issue
Block a user