133 lines
5.0 KiB
Python
133 lines
5.0 KiB
Python
"""命令冷却(CD):给命令触发器加冷却,支持按用户/按群/全局三种作用域
|
|
|
|
基于 rate_limit 的令牌桶实现(rate=1/seconds, capacity=1):
|
|
触发一次取走唯一令牌,seconds 秒内无法再次触发。冷却键由 scope 决定:
|
|
- "user"(默认):按用户 QQ 号 → 同一用户群聊/私聊共用同一 CD(同 nonebot1 传 user_id)
|
|
- "group":按群共享 → 全群一个 CD(同 nonebot1 传 group_id),私聊退化为按用户
|
|
- "global":所有人共享一个 CD
|
|
|
|
用法一(装饰器,包在 on_command 外面):
|
|
from hexi.core.cooldown import cooldown
|
|
|
|
# 单用户 10s CD,群聊/私聊共用
|
|
bf_cmd = cooldown(10, hint="查询太频繁了,{secs} 秒后再试")(on_command("bf3"))
|
|
# 全群共享 60s CD
|
|
daily_cmd = cooldown(60, scope="group")(on_command("打卡"))
|
|
|
|
@bf_cmd.handle()
|
|
async def handler(bot: Bot, ev: MessageEvent):
|
|
...
|
|
|
|
用法二(手动查询/触发,自定义提示文本):
|
|
from hexi.core.cooldown import Cooldown
|
|
|
|
cd = Cooldown(10)
|
|
if cd.in_cd(ev.get_user_id()):
|
|
await matcher.finish(f"冷却中,还剩 {cd.remaining(ev.get_user_id()):.0f} 秒")
|
|
cd.try_trigger(ev.get_user_id())
|
|
"""
|
|
|
|
import itertools
|
|
import math
|
|
from collections.abc import Callable
|
|
from typing import Any, Literal
|
|
|
|
from nonebot.adapters import Event
|
|
from nonebot.dependencies import Dependent
|
|
from nonebot.matcher import Matcher
|
|
|
|
from hexi.core.rate_limit import RateLimiter, get_limiter
|
|
|
|
# 无 name 时的实例序号,保证同一进程内各实例的注册表 key 互不冲突
|
|
_cd_seq = itertools.count(1)
|
|
|
|
Scope = Literal["user", "group", "global"]
|
|
|
|
|
|
def _scope_key(scope: Scope) -> Callable[[Event], str]:
|
|
"""按冷却作用域生成默认冷却键函数(与 key 参数二选一)"""
|
|
if scope == "group":
|
|
|
|
def key(event: Event) -> str:
|
|
# 群聊按群共享;私聊没有群,退化为按用户
|
|
group_id = getattr(event, "group_id", None)
|
|
if group_id is not None:
|
|
return f"group:{group_id}"
|
|
return f"user:{event.get_user_id()}"
|
|
|
|
return key
|
|
if scope == "global":
|
|
return lambda event: "global"
|
|
return lambda event: f"user:{event.get_user_id()}"
|
|
|
|
|
|
class Cooldown:
|
|
"""每用户命令冷却器:同一用户在群聊/私聊共用同一 CD(默认按键为用户 QQ 号)"""
|
|
|
|
def __init__(self, seconds: float, *, name: str | None = None):
|
|
if seconds <= 0:
|
|
raise ValueError(f"seconds 必须大于 0,收到 {seconds!r}")
|
|
self.seconds = float(seconds)
|
|
self.name = name
|
|
self._rate = 1.0 / self.seconds
|
|
|
|
def _limiter(self, key: str) -> RateLimiter:
|
|
# 复用令牌桶:rate=1/seconds、capacity=1 → 取走唯一令牌后 seconds 秒内无法再取
|
|
if self.name is None:
|
|
self.name = f"anon:{next(_cd_seq)}" # 惰性固化,保证 key 稳定
|
|
return get_limiter(f"cd:{self.name}:{key}", rate=self._rate, capacity=1)
|
|
|
|
def try_trigger(self, key: str) -> bool:
|
|
"""尝试触发一次:冷却中返回 False,否则记录触发时刻并返回 True"""
|
|
return self._limiter(key).try_acquire()
|
|
|
|
def in_cd(self, key: str) -> bool:
|
|
"""是否处于冷却中"""
|
|
return self.remaining(key) > 0
|
|
|
|
def remaining(self, key: str) -> float:
|
|
"""剩余冷却秒数(0 表示不在冷却中,可触发)"""
|
|
return self._limiter(key).remaining_secs()
|
|
|
|
|
|
def cooldown(
|
|
seconds: float,
|
|
*,
|
|
hint: str | None = None,
|
|
name: str | None = None,
|
|
scope: Scope = "user",
|
|
key: Callable[[Event], str] | None = None,
|
|
):
|
|
"""给命令触发器(on_command 的返回值)加冷却,返回加好 CD 的触发器
|
|
|
|
- scope: 冷却作用域
|
|
- "user"(默认): 按用户 QQ 号,群聊/私聊共用同一冷却(同 nonebot1 传 user_id)
|
|
- "group": 按群共享(全群一个 CD),私聊退化为按用户
|
|
- "global": 所有人共享一个 CD
|
|
- key: 自定义冷却键函数;传了 key 则忽略 scope
|
|
- hint: 冷却中提示文本,{secs} 会被替换为剩余秒数(向上取整);不传用默认文案
|
|
- name: 冷却器名称;不传时按触发器对象生成,热重载后 CD 会重置
|
|
|
|
原理:向前置插入一个 guard handler,冷却中 finish 提示并终止,
|
|
未冷却则放行后续真正的事件处理函数。
|
|
"""
|
|
cd = Cooldown(seconds, name=name)
|
|
key_fn = key or _scope_key(scope)
|
|
|
|
async def guard(matcher: Matcher, event: Event):
|
|
"""CD 检查 handler:冷却中提示并结束,否则放行"""
|
|
uid = key_fn(event)
|
|
if not cd.try_trigger(uid):
|
|
secs = math.ceil(cd.remaining(uid))
|
|
msg = (hint or "操作太频繁了,请 {secs} 秒后再试").format(secs=secs)
|
|
await matcher.finish(msg)
|
|
|
|
def deco(matcher: type[Matcher]) -> type[Matcher]:
|
|
if cd.name is None:
|
|
cd.name = f"cd:{next(_cd_seq)}"
|
|
dep = Dependent[Any].parse(call=guard, allow_types=Matcher.HANDLER_PARAM_TYPES)
|
|
matcher.handlers.insert(0, dep)
|
|
return matcher
|
|
|
|
return deco
|