From 31621768c52c85db3e1db86262ffa00f7ab71d4c Mon Sep 17 00:00:00 2001 From: Misty02600 Date: Thu, 13 Aug 2026 10:51:02 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BD=BF=E7=94=A8=E8=B7=A8=E5=B9=B3?= =?UTF-8?q?=E5=8F=B0=E5=AE=9E=E4=BD=93=E9=94=AE=E6=8C=81=E4=B9=85=E5=8C=96?= =?UTF-8?q?=E7=94=A8=E9=87=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nonebot_plugin_limiter/cooldown.py | 29 ++++---- nonebot_plugin_limiter/entity.py | 37 +++++++---- nonebot_plugin_limiter/entity_key.py | 98 ++++++++++++++++++++++++++++ nonebot_plugin_limiter/persist.py | 13 ++-- 4 files changed, 146 insertions(+), 31 deletions(-) create mode 100644 nonebot_plugin_limiter/entity_key.py diff --git a/nonebot_plugin_limiter/cooldown.py b/nonebot_plugin_limiter/cooldown.py index 8f35bd8..79e8350 100644 --- a/nonebot_plugin_limiter/cooldown.py +++ b/nonebot_plugin_limiter/cooldown.py @@ -20,12 +20,15 @@ from tzlocal import get_localzone from .entity import BYPASS_ENTITY, CooldownEntity +from .entity_key import EntityKey _tz = get_localzone() SupportMsgType = str | Message | MessageSegment | MessageTemplate | UniMessage -def _entity_id_dep_wrapper(entity: CooldownEntity | _DependentCallable[str]) -> _DependentCallable[str]: +def _entity_id_dep_wrapper( + entity: CooldownEntity | _DependentCallable[EntityKey], +) -> _DependentCallable[EntityKey]: if isinstance(entity, CooldownEntity): entity_id_dep = entity.get_entity_id else: @@ -85,11 +88,11 @@ class FixWindowUsage: available: int -_FixWindowCooldownDict: dict[str, dict[str, FixWindowUsage]] = {} +_FixWindowCooldownDict: dict[str, dict[EntityKey, FixWindowUsage]] = {} def Cooldown( - entity: CooldownEntity | _DependentCallable[str], + entity: CooldownEntity | _DependentCallable[EntityKey], period: int | timedelta | str, *, limit: int | _DependentCallable[int] = 5, @@ -163,11 +166,11 @@ async def handler(...): ... _FixWindowCooldownDict[name] = {} bucket = _FixWindowCooldownDict[name] else: - bucket: dict[str, FixWindowUsage] = {} + bucket: dict[EntityKey, FixWindowUsage] = {} async def _limiter_dependency( state: T_State, - entity_id: str = Depends(_entity_id_dep_wrapper(entity)), + entity_id: EntityKey = Depends(_entity_id_dep_wrapper(entity)), limit: int = Depends(_limit_dep_wrapper(limit)), reject_cb: Callable[..., Awaitable[Any]] = Depends(_reject_dep_wrapper(reject)), ) -> None: @@ -215,11 +218,11 @@ class SlidingWindowUsage: timestamps: deque[datetime] = field(default_factory=deque) -_SlidingWindowCooldownDict: dict[str, dict[str, SlidingWindowUsage]] = {} +_SlidingWindowCooldownDict: dict[str, dict[EntityKey, SlidingWindowUsage]] = {} def SlidingWindowCooldown( - entity: CooldownEntity | _DependentCallable[str], + entity: CooldownEntity | _DependentCallable[EntityKey], period: int | timedelta, *, limit: int | _DependentCallable[int] = 5, @@ -284,11 +287,11 @@ async def handler(...): ... if isinstance(name, str): bucket = _SlidingWindowCooldownDict.setdefault(name, {}) else: - bucket: dict[str, SlidingWindowUsage] = {} + bucket: dict[EntityKey, SlidingWindowUsage] = {} async def _limiter_dependency( state: T_State, - entity_id: str = Depends(_entity_id_dep_wrapper(entity)), + entity_id: EntityKey = Depends(_entity_id_dep_wrapper(entity)), limit: int = Depends(_limit_dep_wrapper(limit)), reject_cb: Callable[..., Awaitable[Any]] = Depends(_reject_dep_wrapper(reject)), ) -> None: @@ -333,11 +336,11 @@ class TokenBucketUsage: available: int -_TokenBucketCooldownDict: dict[str, dict[str, TokenBucketUsage]] = {} +_TokenBucketCooldownDict: dict[str, dict[EntityKey, TokenBucketUsage]] = {} def TokenBucketCooldown( - entity: CooldownEntity | _DependentCallable[str], + entity: CooldownEntity | _DependentCallable[EntityKey], capacity: int, add_speed: int, *, @@ -404,11 +407,11 @@ async def handler(...): ... _TokenBucketCooldownDict[name] = {} bucket = _TokenBucketCooldownDict[name] else: - bucket: dict[str, TokenBucketUsage] = {} + bucket: dict[EntityKey, TokenBucketUsage] = {} async def _limiter_dependency( state: T_State, - entity_id: str = Depends(_entity_id_dep_wrapper(entity)), + entity_id: EntityKey = Depends(_entity_id_dep_wrapper(entity)), consume_size: int = Depends(_limit_dep_wrapper(consume_size)), reject_cb: Callable[..., Awaitable[Any]] = Depends(_reject_dep_wrapper(reject)), ) -> None: diff --git a/nonebot_plugin_limiter/entity.py b/nonebot_plugin_limiter/entity.py index 3f9f07d..a3e0d5a 100644 --- a/nonebot_plugin_limiter/entity.py +++ b/nonebot_plugin_limiter/entity.py @@ -5,6 +5,8 @@ from nonebot.permission import Permission from nonebot_plugin_uninfo import get_session +from .entity_key import EntityKey, SceneUserKey, make_scene_key, make_user_key + _IdType = str | int BYPASS_ENTITY = "__bypass" @@ -20,7 +22,7 @@ class CooldownEntity: def __init__(self) -> None: ... @abstractmethod - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: """ 返回被限制实体的唯一标识符，统一为 str """ @@ -37,7 +39,7 @@ class GlobalScope(CooldownEntity): def __init__(self) -> None: pass - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: return "__global" @@ -67,7 +69,7 @@ def __init__(self, *, whitelist: None | tuple[_IdType, ...] = None, permission: self.whitelist = None self.permission = permission - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: sess = await get_session(bot, event) if sess is None: return BYPASS_ENTITY @@ -77,7 +79,7 @@ async def get_entity_id(self, bot: Bot, event: Event) -> str: return BYPASS_ENTITY if self.permission is not None and (await self.permission(bot, event)): return BYPASS_ENTITY - return f"u`{user_id}`" + return make_user_key(sess.scope, user_id) class SceneScope(CooldownEntity): @@ -104,7 +106,7 @@ def __init__(self, *, whitelist: None | tuple[_IdType, ...] = None, permission: self.whitelist = None self.permission = permission - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: sess = await get_session(bot, event) if sess is None: return BYPASS_ENTITY @@ -114,7 +116,12 @@ async def get_entity_id(self, bot: Bot, event: Event) -> str: return BYPASS_ENTITY if self.permission is not None and (await self.permission(bot, event)): return BYPASS_ENTITY - return f"s`{scene_id}`" + return make_scene_key( + sess.scope, + scene_id, + sess.scene.type, + sess.scene.parent.id if sess.scene.parent else None, + ) class UserSceneScope(CooldownEntity): @@ -149,7 +156,7 @@ def __init__( self.whitelist = None self.permission = permission - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: sess = await get_session(bot, event) if sess is None: return BYPASS_ENTITY @@ -162,7 +169,13 @@ async def get_entity_id(self, bot: Bot, event: Event) -> str: return BYPASS_ENTITY if self.permission is not None and (await self.permission(bot, event)): return BYPASS_ENTITY - return f"u`{user_id}`_s`{scene_id}`" + scene_key = make_scene_key( + sess.scope, + scene_id, + sess.scene.type, + sess.scene.parent.id if sess.scene.parent else None, + ) + return SceneUserKey(scene_key, user_id) class PrivateScope(CooldownEntity): @@ -191,7 +204,7 @@ def __init__(self, *, whitelist: None | tuple[_IdType, ...] = None, permission: self.whitelist = None self.permission = permission - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: sess = await get_session(bot, event) if sess is None or not sess.scene.is_private: return BYPASS_ENTITY @@ -201,7 +214,7 @@ async def get_entity_id(self, bot: Bot, event: Event) -> str: return BYPASS_ENTITY if self.permission is not None and (await self.permission(bot, event)): return BYPASS_ENTITY - return f"u`{user_id}`" + return make_user_key(sess.scope, user_id) class PublicScope(CooldownEntity): @@ -230,7 +243,7 @@ def __init__(self, *, whitelist: None | tuple[_IdType, ...] = None, permission: self.whitelist = None self.permission = permission - async def get_entity_id(self, bot: Bot, event: Event) -> str: + async def get_entity_id(self, bot: Bot, event: Event) -> EntityKey: sess = await get_session(bot, event) if sess is None or sess.scene.is_private: return BYPASS_ENTITY @@ -240,4 +253,4 @@ async def get_entity_id(self, bot: Bot, event: Event) -> str: return BYPASS_ENTITY if self.permission is not None and (await self.permission(bot, event)): return BYPASS_ENTITY - return f"u`{user_id}`" + return make_user_key(sess.scope, user_id) diff --git a/nonebot_plugin_limiter/entity_key.py b/nonebot_plugin_limiter/entity_key.py new file mode 100644 index 0000000..ab7b3d6 --- /dev/null +++ b/nonebot_plugin_limiter/entity_key.py @@ -0,0 +1,98 @@ +from dataclasses import dataclass +import json +from typing import TypeAlias + +from nonebot_plugin_uninfo import SceneType, SupportScope + +ENTITY_KEY_PREFIX = "limiter-entity:" + + +@dataclass(frozen=True) +class UserKey: + scope: str + user_id: str + + +@dataclass(frozen=True) +class SceneKey: + scope: str + scene_id: str + scene_type: int + parent_id: str | None = None + + +@dataclass(frozen=True) +class SceneUserKey: + scene: SceneKey + user_id: str + + +EntityKey: TypeAlias = str | UserKey | SceneKey | SceneUserKey + + +def make_user_key(scope: SupportScope, user_id: str) -> UserKey: + return UserKey(scope.value, user_id) + + +def make_scene_key( + scope: SupportScope, + scene_id: str, + scene_type: SceneType, + parent_id: str | None, +) -> SceneKey: + return SceneKey(scope.value, scene_id, scene_type.value, parent_id) + + +def serialize_entity_key(key: EntityKey) -> str: + match key: + case str(): + return key + case UserKey(scope, user_id): + value = ["user", scope, user_id] + case SceneKey(scope, scene_id, scene_type, parent_id): + value = ["scene", scope, scene_id, scene_type, parent_id] + case SceneUserKey(SceneKey(scope, scene_id, scene_type, parent_id), user_id): + value = [ + "scene_user", + scope, + scene_id, + scene_type, + parent_id, + user_id, + ] + case _: + raise TypeError(f"unsupported entity key type: {type(key)!r}") + payload = json.dumps(value, ensure_ascii=False, separators=(",", ":")) + return f"{ENTITY_KEY_PREFIX}{payload}" + + +def deserialize_entity_key(value: str) -> EntityKey: + if not value.startswith(ENTITY_KEY_PREFIX): + return value + + try: + data = json.loads(value[len(ENTITY_KEY_PREFIX) :]) + if not isinstance(data, list) or not data or not isinstance(data[0], str): + return value + kind = data[0] + if kind == "user" and len(data) == 3 and all(isinstance(item, str) for item in data[1:]): + return UserKey(data[1], data[2]) + if kind == "scene" and _is_scene_payload(data): + return SceneKey(data[1], data[2], data[3], data[4]) + if kind == "scene_user" and len(data) == 6 and _is_scene_payload(data[:5]) and isinstance(data[5], str): + scene = SceneKey(data[1], data[2], data[3], data[4]) + return SceneUserKey(scene, data[5]) + except json.JSONDecodeError: + pass + return value + + +def _is_scene_payload(value: list[object]) -> bool: + return ( + len(value) == 5 + and isinstance(value[1], str) + and isinstance(value[2], str) + and isinstance(value[3], int) + and not isinstance(value[3], bool) + and (value[4] is None or isinstance(value[4], str)) + ) diff --git a/nonebot_plugin_limiter/persist.py b/nonebot_plugin_limiter/persist.py index 5d122f2..4a08be0 100644 --- a/nonebot_plugin_limiter/persist.py +++ b/nonebot_plugin_limiter/persist.py @@ -20,6 +20,7 @@ _TokenBucketCooldownDict, _tz, ) +from .entity_key import deserialize_entity_key, serialize_entity_key driver = get_driver() plugin_data_file: Path = store.get_plugin_data_file("limiter_data.json") @@ -62,7 +63,7 @@ def load_usage_data() -> None: _FixWindowCooldownDict[name] = {} bucket = _FixWindowCooldownDict[name] for _id, usage in usage_set.items(): - bucket[_id] = FixWindowUsage( + bucket[deserialize_entity_key(_id)] = FixWindowUsage( start_time=datetime.fromtimestamp(usage.start_time, tz=_tz), available=usage.available, ) @@ -73,7 +74,7 @@ def load_usage_data() -> None: _SlidingWindowCooldownDict[name] = {} bucket = _SlidingWindowCooldownDict[name] for _id, usage in usage_set.items(): - bucket[_id] = SlidingWindowUsage( + bucket[deserialize_entity_key(_id)] = SlidingWindowUsage( timestamps=deque(datetime.fromtimestamp(t, tz=_tz) for t in usage.timestamps) ) @@ -83,7 +84,7 @@ def load_usage_data() -> None: _TokenBucketCooldownDict[name] = {} bucket = _TokenBucketCooldownDict[name] for _id, usage in usage_set.items(): - bucket[_id] = TokenBucketUsage( + bucket[deserialize_entity_key(_id)] = TokenBucketUsage( last_update_time=datetime.fromtimestamp(usage.last_update_time, tz=_tz), capacity=usage.capacity, available=usage.available, @@ -99,7 +100,7 @@ def save_usage_data() -> None: for name, usage_set in _FixWindowCooldownDict.items(): j["fix_window"][name] = { - _id: { + serialize_entity_key(_id): { "start_time": int(usage.start_time.timestamp()), "available": usage.available, } @@ -108,7 +109,7 @@ def save_usage_data() -> None: for name, usage_set in _SlidingWindowCooldownDict.items(): j["sliding_window"][name] = { - _id: { + serialize_entity_key(_id): { "timestamps": [int(t.timestamp()) for t in usage.timestamps], } for _id, usage in usage_set.items() @@ -116,7 +117,7 @@ def save_usage_data() -> None: for name, usage_set in _TokenBucketCooldownDict.items(): j["token_bucket"][name] = { - _id: { + serialize_entity_key(_id): { "last_update_time": int(usage.last_update_time.timestamp()), "capacity": usage.capacity, "available": usage.available,