Feature: Add Increaser

Author: middlered

README.md

@@ -138,6 +138,22 @@ cmd2 = on_startswith("cmd2")
 async def _(): ...
 ```
 
+手动管理限制器使用记录的增加操作，适用于没有成功完成事件处理时避免给用户添加限制。
+```python
+from nonebot_plugin_limiter import Increaser
+
+cmd = on_startswith("cmd")
+@cmd.handle(parameterless=[
+    Cooldown(UserScope(), 100, limit = 2, set_increaser = True)
+])
+async def _(increaser: Increaser):
+    if <not_meet_condition>:
+        await cmd.finish("Run failed")
+    else:
+        increaser.execute() # 这会给该实体添加一条使用记录
+        await cmd.finish("Run successed")
+```
+
 使用固定窗口策略实现每日签到
 ```python
 dailysign = on_startswith("签到")

nonebot_plugin_limiter/__init__.py

@@ -28,5 +28,6 @@
 from .entity import SceneScope as SceneScope
 from .entity import UserSceneScope as UserSceneScope
 from .entity import UserScope as UserScope
+from .handler import Increaser as Increaser
 from .persist import load_usage_data as load_usage_data
 from .persist import save_usage_data as save_usage_data

nonebot_plugin_limiter/cooldown.py

@@ -10,7 +10,7 @@
 from nonebot.matcher import Matcher
 from nonebot.params import Depends
 from nonebot.rule import Rule as Rule
-from nonebot.typing import _DependentCallable
+from nonebot.typing import T_State, _DependentCallable
 from nonebot_plugin_alconna import UniMessage
 from tzlocal import get_localzone
 
@@ -33,6 +33,7 @@ def _limit_dep_wrapper(limit: int | _DependentCallable[int]) -> _DependentCallab
         limit_dep = limit
     return limit_dep
 
+# region: FixWindow
 @dataclass
 class FixWindowUsage:
     start_time: datetime
@@ -48,6 +49,7 @@ def Cooldown(
     *,
     limit: int | _DependentCallable[int] = 5,
     reject: None | str | Message | MessageSegment | MessageTemplate | UniMessage = None,
+    set_increaser: bool = False,
     name: None | str = None,
 ):
     """
@@ -120,6 +122,7 @@ async def _limiter_dependency(
         bot: Bot,
         matcher: Matcher,
         event: Event,
+        state: T_State,
         entity_id: str = Depends(entity_id_dep),
         limit: int = Depends(limit_dep),
     ) -> None:
@@ -138,21 +141,30 @@ async def _limiter_dependency(
 
         # Calculate reset time based on when the limitation was set
         reset_time = trigger.get_next_fire_time(usage.start_time, now)
-
         assert reset_time is not None, "reset_time should not be None"
-        if now >= reset_time:
+
+        def _increase_action():
             usage.start_time = now
             usage.available = limit - 1
-            return
 
+        if now >= reset_time:
+            if set_increaser:
+                state["plugin_limiter:increaser"] = _increase_action
+            else:
+                _increase_action()
+            return  # Didn't exceed
+
+        # Exceed
         if isinstance(reject, UniMessage):
             await reject.finish(event, bot)
         else:
             await matcher.finish(reject)
 
     return Depends(_limiter_dependency)
 
+# endregin
 
+# region: SlidingWindow
 @dataclass
 class SlidingWindowUsage:
     timestamps: deque[datetime] = field(default_factory=deque)
@@ -167,6 +179,7 @@ def SlidingWindowCooldown(
     *,
     limit: int | _DependentCallable[int] = 5,
     reject: None | str | Message | MessageSegment | MessageTemplate | UniMessage = None,
+    set_increaser: bool = False,
     name: None | str = None,
 ):
     """
@@ -191,6 +204,10 @@ def SlidingWindowCooldown(
             可选，当超出限制时的响应行为。默认为 `None`。
             - 若为 `str` 或消息对象，将作为限制使用时的提示消息发送给用户。
 
+        set_increaser (bool):
+            可选，是否获取限制器的增加器。默认为 False。
+            - 当启用该选项时，限制器默认的自增将会关闭，需要在事件处理时依赖获取 Increaser 并手动操作增加。
+
         name (None | str):
             可选，设置当前限制器的使用统计集合。默认为 `None` ，即私有集合。
             - 当传入 `str` ，将创建或加入一个同名公共集合，可用于与其他命令的限制器共享使用统计。
@@ -230,6 +247,7 @@ async def _limiter_dependency(
         bot: Bot,
         matcher: Matcher,
         event: Event,
+        state: T_State,
         entity_id: str = Depends(entity_id_dep),
         limit: int = Depends(limit_dep),
     ) -> None:
@@ -246,13 +264,22 @@ async def _limiter_dependency(
         while usage.timestamps and (now - usage.timestamps[0]).total_seconds() >= window_length:
             usage.timestamps.popleft()
 
-        if len(usage.timestamps) < limit:
+        def _increase_action():
             usage.timestamps.append(now)
-            return
 
+        if len(usage.timestamps) < limit:
+            if set_increaser:
+                state["plugin_limiter:increaser"] = _increase_action
+            else:
+                _increase_action()
+            return  # Didn't exceed
+
+        # Exceeded
         if isinstance(reject, UniMessage):
             await reject.finish(event, bot)
         else:
             await matcher.finish(reject)
 
     return Depends(_limiter_dependency)
+
+# endregin

nonebot_plugin_limiter/handler.py

@@ -1,5 +1,21 @@
+from collections.abc import Callable
+from typing import Annotated
+
 from nonebot.params import Depends
-from nonebot.typing import T_State, _DependentCallable
+from nonebot.typing import T_State
+
+
+class _FuncWrapper:
+    def __init__(self, func: Callable) -> None:
+        self._func = func
+
+    def execute(self):
+        self._func()
+
+def get_increaser(state: T_State):
+    ret = state.get("plugin_limiter:increaser")
+    if ret is None:
+        raise KeyError("Cannot get increaser, make sure you have enabled `set_increaser` in cooldown policy.")
+    return _FuncWrapper(ret)
 
-class Limiter:
-    pass
+Increaser = Annotated[_FuncWrapper, Depends(get_increaser)]