Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 16 additions & 13 deletions nonebot_plugin_limiter/cooldown.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
*,
Expand Down Expand Up @@ -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:
Expand Down
37 changes: 25 additions & 12 deletions nonebot_plugin_limiter/entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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
"""
Expand All @@ -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"


Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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)
98 changes: 98 additions & 0 deletions nonebot_plugin_limiter/entity_key.py
Original file line number Diff line number Diff line change
@@ -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:"

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

正常的 key 我觉得完全不需要 prefix,也不存别的什么地方



@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:

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

没看懂这么写包了层为了啥,直接实例化不行么

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

而且其实我目目前目测看并没有看到需要用 dataclass 的必要性, NamedTuple 完全可以解决

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

你说的对( 我没仔细检查

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))
)
13 changes: 7 additions & 6 deletions nonebot_plugin_limiter/persist.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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,
)
Expand All @@ -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)
)

Expand All @@ -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,
Expand All @@ -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,
}
Expand All @@ -108,15 +109,15 @@ 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()
}

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,
Expand Down