hikari-arc 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- arc/__init__.py +8 -2
- arc/abc/__init__.py +10 -1
- arc/abc/client.py +116 -30
- arc/abc/command.py +54 -10
- arc/abc/error_handler.py +10 -2
- arc/abc/limiter.py +5 -16
- arc/abc/option.py +71 -3
- arc/abc/plugin.py +92 -0
- arc/client.py +3 -3
- arc/command/message.py +4 -4
- arc/command/option/attachment.py +3 -3
- arc/command/option/bool.py +3 -4
- arc/command/option/channel.py +3 -3
- arc/command/option/float.py +4 -4
- arc/command/option/int.py +4 -4
- arc/command/option/mentionable.py +3 -3
- arc/command/option/role.py +3 -3
- arc/command/option/str.py +4 -4
- arc/command/option/user.py +3 -3
- arc/command/slash.py +116 -115
- arc/command/user.py +3 -3
- arc/context/base.py +114 -20
- arc/errors.py +31 -12
- arc/events.py +5 -1
- arc/internal/about.py +1 -1
- arc/internal/options.py +106 -0
- arc/internal/sigparse.py +19 -1
- arc/internal/sync.py +7 -2
- arc/internal/types.py +7 -12
- arc/plugin.py +2 -2
- arc/utils/__init__.py +4 -1
- arc/utils/hooks/__init__.py +2 -2
- arc/utils/hooks/limiters.py +70 -202
- arc/utils/ratelimiter.py +243 -0
- {hikari_arc-0.5.0.dist-info → hikari_arc-0.6.0.dist-info}/METADATA +11 -6
- hikari_arc-0.6.0.dist-info/RECORD +52 -0
- hikari_arc-0.5.0.dist-info/RECORD +0 -50
- {hikari_arc-0.5.0.dist-info → hikari_arc-0.6.0.dist-info}/LICENSE +0 -0
- {hikari_arc-0.5.0.dist-info → hikari_arc-0.6.0.dist-info}/WHEEL +0 -0
- {hikari_arc-0.5.0.dist-info → hikari_arc-0.6.0.dist-info}/top_level.txt +0 -0
arc/utils/hooks/limiters.py
CHANGED
|
@@ -1,25 +1,15 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
import asyncio
|
|
4
|
-
import sys
|
|
5
|
-
import time
|
|
6
|
-
import traceback
|
|
7
3
|
import typing as t
|
|
8
|
-
from collections import deque
|
|
9
4
|
|
|
10
|
-
import attr
|
|
11
|
-
|
|
12
|
-
from arc.abc.hookable import HookResult
|
|
13
5
|
from arc.abc.limiter import LimiterProto
|
|
6
|
+
from arc.context.base import Context
|
|
14
7
|
from arc.errors import UnderCooldownError
|
|
15
8
|
from arc.internal.types import ClientT
|
|
16
|
-
|
|
17
|
-
if t.TYPE_CHECKING:
|
|
18
|
-
from arc.context.base import Context
|
|
19
|
-
|
|
9
|
+
from arc.utils.ratelimiter import RateLimiter, RateLimiterExhaustedError
|
|
20
10
|
|
|
21
11
|
__all__ = (
|
|
22
|
-
"
|
|
12
|
+
"LimiterHook",
|
|
23
13
|
"global_limiter",
|
|
24
14
|
"guild_limiter",
|
|
25
15
|
"channel_limiter",
|
|
@@ -29,89 +19,8 @@ __all__ = (
|
|
|
29
19
|
)
|
|
30
20
|
|
|
31
21
|
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
"""Handles the ratelimiting of a single item. (E.g. a single user or a channel)."""
|
|
35
|
-
|
|
36
|
-
key: str
|
|
37
|
-
"""The key of the bucket."""
|
|
38
|
-
|
|
39
|
-
reset_at: float
|
|
40
|
-
"""The time at which the bucket resets."""
|
|
41
|
-
|
|
42
|
-
limiter: RateLimiter[ClientT]
|
|
43
|
-
"""The limiter this bucket belongs to."""
|
|
44
|
-
|
|
45
|
-
_remaining: int = attr.field(alias="remaining")
|
|
46
|
-
"""The amount of requests remaining until the bucket is exhausted."""
|
|
47
|
-
|
|
48
|
-
_queue: deque[asyncio.Event] = attr.field(factory=deque, init=False)
|
|
49
|
-
"""A list of events to set as the iter task proceeds."""
|
|
50
|
-
|
|
51
|
-
_task: asyncio.Task[None] | None = attr.field(default=None, init=False)
|
|
52
|
-
"""The task that is currently iterating over the queue."""
|
|
53
|
-
|
|
54
|
-
@classmethod
|
|
55
|
-
def for_limiter(cls, key: str, limiter: RateLimiter[ClientT]) -> _Bucket[ClientT]:
|
|
56
|
-
"""Create a new bucket for a RateLimiter."""
|
|
57
|
-
return cls(key=key, limiter=limiter, reset_at=time.monotonic() + limiter.period, remaining=limiter.limit)
|
|
58
|
-
|
|
59
|
-
@property
|
|
60
|
-
def remaining(self) -> int:
|
|
61
|
-
"""The amount of requests remaining until the bucket is exhausted."""
|
|
62
|
-
if self.reset_at <= time.monotonic():
|
|
63
|
-
self.reset()
|
|
64
|
-
return self._remaining
|
|
65
|
-
|
|
66
|
-
@remaining.setter
|
|
67
|
-
def remaining(self, value: int) -> None:
|
|
68
|
-
self._remaining = value
|
|
69
|
-
|
|
70
|
-
@property
|
|
71
|
-
def is_exhausted(self) -> bool:
|
|
72
|
-
"""Return a boolean determining if the bucket is exhausted."""
|
|
73
|
-
return self.remaining <= 0 and self.reset_at > time.monotonic()
|
|
74
|
-
|
|
75
|
-
@property
|
|
76
|
-
def is_stale(self) -> bool:
|
|
77
|
-
"""Return a boolean determining if the bucket is stale.
|
|
78
|
-
If a bucket is stale, it is no longer in use and can be purged.
|
|
79
|
-
"""
|
|
80
|
-
return not self._queue and self.remaining == self.limiter.limit and (self._task is None or self._task.done())
|
|
81
|
-
|
|
82
|
-
def start_queue(self) -> None:
|
|
83
|
-
"""Start the queue of a bucket.
|
|
84
|
-
This will start setting events in the queue until the bucket is ratelimited.
|
|
85
|
-
"""
|
|
86
|
-
if self._task is None or self._task.done():
|
|
87
|
-
self._task = asyncio.create_task(self._iter_queue())
|
|
88
|
-
|
|
89
|
-
def reset(self) -> None:
|
|
90
|
-
"""Reset the bucket."""
|
|
91
|
-
self.reset_at = time.monotonic() + self.limiter.period
|
|
92
|
-
self._remaining = self.limiter.limit
|
|
93
|
-
|
|
94
|
-
async def _iter_queue(self) -> None:
|
|
95
|
-
"""Iterate over the queue and set events until exhausted."""
|
|
96
|
-
try:
|
|
97
|
-
while self._queue:
|
|
98
|
-
if self.remaining <= 0 and self.reset_at > time.monotonic():
|
|
99
|
-
# Sleep until ratelimit expires
|
|
100
|
-
await asyncio.sleep(self.reset_at - time.monotonic())
|
|
101
|
-
self.reset()
|
|
102
|
-
|
|
103
|
-
# Set events while not ratelimited
|
|
104
|
-
while self.remaining > 0 and self._queue:
|
|
105
|
-
self._queue.popleft().set()
|
|
106
|
-
self._remaining -= 1
|
|
107
|
-
|
|
108
|
-
except Exception as e:
|
|
109
|
-
print(f"Task Exception was never retrieved: {e}", file=sys.stderr)
|
|
110
|
-
print(traceback.format_exc(), file=sys.stderr)
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
class RateLimiter(LimiterProto[ClientT]):
|
|
114
|
-
"""The default implementation of a ratelimiter.
|
|
22
|
+
class LimiterHook(RateLimiter[Context[ClientT]], LimiterProto[ClientT]):
|
|
23
|
+
"""The default implementation of a ratelimiter that can be used as a hook.
|
|
115
24
|
|
|
116
25
|
Parameters
|
|
117
26
|
----------
|
|
@@ -119,58 +28,20 @@ class RateLimiter(LimiterProto[ClientT]):
|
|
|
119
28
|
The period, in seconds, after which the bucket resets.
|
|
120
29
|
limit : int
|
|
121
30
|
The amount of requests allowed in a bucket.
|
|
122
|
-
get_key_with : Callable[[Context[t.Any]], str]
|
|
31
|
+
get_key_with : t.Callable[[Context[t.Any]], str]
|
|
123
32
|
A callable that returns a key for the ratelimiter bucket.
|
|
124
|
-
"""
|
|
125
|
-
|
|
126
|
-
__slots__ = ("period", "limit", "_buckets", "_get_key")
|
|
127
|
-
|
|
128
|
-
def __init__(self, period: float, limit: int, *, get_key_with: t.Callable[[Context[t.Any]], str]) -> None:
|
|
129
|
-
self.period: float = period
|
|
130
|
-
self.limit: int = limit
|
|
131
|
-
self._buckets: t.Dict[str, _Bucket[ClientT]] = {}
|
|
132
|
-
self._get_key: t.Callable[[Context[t.Any]], str] = get_key_with
|
|
133
|
-
self._gc_task: asyncio.Task[None] | None = None
|
|
134
33
|
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
ctx : Context[t.Any]
|
|
145
|
-
The context to evaluate the ratelimit under.
|
|
34
|
+
See Also
|
|
35
|
+
--------
|
|
36
|
+
- [`global_limiter`][arc.utils.hooks.limiters.global_limiter]
|
|
37
|
+
- [`guild_limiter`][arc.utils.hooks.limiters.guild_limiter]
|
|
38
|
+
- [`channel_limiter`][arc.utils.hooks.limiters.channel_limiter]
|
|
39
|
+
- [`user_limiter`][arc.utils.hooks.limiters.user_limiter]
|
|
40
|
+
- [`member_limiter`][arc.utils.hooks.limiters.member_limiter]
|
|
41
|
+
- [`custom_limiter`][arc.utils.hooks.limiters.custom_limiter]
|
|
42
|
+
"""
|
|
146
43
|
|
|
147
|
-
|
|
148
|
-
-------
|
|
149
|
-
bool
|
|
150
|
-
A boolean determining if the ratelimiter is ratelimited or not.
|
|
151
|
-
"""
|
|
152
|
-
now = time.monotonic()
|
|
153
|
-
|
|
154
|
-
if data := self._buckets.get(self.get_key(ctx)):
|
|
155
|
-
if data.reset_at <= now:
|
|
156
|
-
return False
|
|
157
|
-
return data._remaining <= 0
|
|
158
|
-
return False
|
|
159
|
-
|
|
160
|
-
def _start_gc(self) -> None:
|
|
161
|
-
"""Start the garbage collector task if one is not running."""
|
|
162
|
-
if self._gc_task is None or self._gc_task.done():
|
|
163
|
-
self._gc_task = asyncio.create_task(self._gc())
|
|
164
|
-
|
|
165
|
-
async def _gc(self) -> None:
|
|
166
|
-
"""Purge stale buckets."""
|
|
167
|
-
while self._buckets:
|
|
168
|
-
await asyncio.sleep(self.period + 1.0)
|
|
169
|
-
for bucket in list(self._buckets.values()):
|
|
170
|
-
if bucket.is_stale:
|
|
171
|
-
del self._buckets[bucket.key]
|
|
172
|
-
|
|
173
|
-
async def acquire(self, ctx: Context[t.Any], *, wait: bool = True) -> None:
|
|
44
|
+
async def acquire(self, ctx: Context[ClientT], *, wait: bool = True) -> None:
|
|
174
45
|
"""Acquire a bucket, block execution if ratelimited and wait is True.
|
|
175
46
|
|
|
176
47
|
Parameters
|
|
@@ -186,54 +57,15 @@ class RateLimiter(LimiterProto[ClientT]):
|
|
|
186
57
|
UnderCooldownError
|
|
187
58
|
If the ratelimiter is ratelimited and wait is False.
|
|
188
59
|
"""
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
# Get existing or insert new bucket
|
|
193
|
-
bucket = self._buckets.setdefault(key, _Bucket.for_limiter(key, self))
|
|
194
|
-
|
|
195
|
-
if bucket.is_exhausted and not wait:
|
|
60
|
+
try:
|
|
61
|
+
return await super().acquire(ctx, wait=wait)
|
|
62
|
+
except RateLimiterExhaustedError as exc:
|
|
196
63
|
raise UnderCooldownError(
|
|
197
|
-
self,
|
|
198
|
-
|
|
199
|
-
f"Ratelimited for {bucket.reset_at - time.monotonic()} seconds.",
|
|
200
|
-
)
|
|
64
|
+
self, exc.retry_after, f"Command is under cooldown for '{exc.retry_after}' seconds."
|
|
65
|
+
) from exc
|
|
201
66
|
|
|
202
|
-
bucket._queue.append(event)
|
|
203
|
-
bucket.start_queue()
|
|
204
|
-
self._start_gc()
|
|
205
|
-
|
|
206
|
-
if wait:
|
|
207
|
-
await event.wait()
|
|
208
|
-
|
|
209
|
-
async def __call__(self, ctx: Context[t.Any]) -> HookResult:
|
|
210
|
-
"""Acquire a ratelimit, fail if ratelimited.
|
|
211
|
-
|
|
212
|
-
Parameters
|
|
213
|
-
----------
|
|
214
|
-
ctx : Context[t.Any]
|
|
215
|
-
The context to evaluate the ratelimit under.
|
|
216
67
|
|
|
217
|
-
|
|
218
|
-
-------
|
|
219
|
-
HookResult
|
|
220
|
-
A hook result to conform to the limiter protocol.
|
|
221
|
-
|
|
222
|
-
Raises
|
|
223
|
-
------
|
|
224
|
-
UnderCooldownError
|
|
225
|
-
If the ratelimiter is ratelimited.
|
|
226
|
-
"""
|
|
227
|
-
await self.acquire(ctx, wait=False)
|
|
228
|
-
return HookResult()
|
|
229
|
-
|
|
230
|
-
def reset(self, ctx: Context[t.Any]) -> None:
|
|
231
|
-
"""Reset the ratelimit for a given context."""
|
|
232
|
-
if bucket := self._buckets.get(self.get_key(ctx)):
|
|
233
|
-
bucket.reset()
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
def global_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
68
|
+
def global_limiter(period: float, limit: int) -> LimiterHook[t.Any]:
|
|
237
69
|
"""Create a global ratelimiter.
|
|
238
70
|
|
|
239
71
|
This ratelimiter is shared across all invocation contexts.
|
|
@@ -244,11 +76,17 @@ def global_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
|
244
76
|
The period, in seconds, after which the bucket resets.
|
|
245
77
|
limit : int
|
|
246
78
|
The amount of requests allowed in a bucket.
|
|
79
|
+
|
|
80
|
+
Usage
|
|
81
|
+
-----
|
|
82
|
+
```py
|
|
83
|
+
@arc.with_hook(arc.global_limiter(5.0, 1))
|
|
84
|
+
```
|
|
247
85
|
"""
|
|
248
|
-
return
|
|
86
|
+
return LimiterHook(period, limit, get_key_with=lambda _: "amongus")
|
|
249
87
|
|
|
250
88
|
|
|
251
|
-
def guild_limiter(period: float, limit: int) ->
|
|
89
|
+
def guild_limiter(period: float, limit: int) -> LimiterHook[t.Any]:
|
|
252
90
|
"""Create a guild ratelimiter.
|
|
253
91
|
|
|
254
92
|
This ratelimiter is shared across all contexts in a guild.
|
|
@@ -259,11 +97,17 @@ def guild_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
|
259
97
|
The period, in seconds, after which the bucket resets.
|
|
260
98
|
limit : int
|
|
261
99
|
The amount of requests allowed in a bucket.
|
|
100
|
+
|
|
101
|
+
Usage
|
|
102
|
+
-----
|
|
103
|
+
```py
|
|
104
|
+
@arc.with_hook(arc.guild_limiter(5.0, 1))
|
|
105
|
+
```
|
|
262
106
|
"""
|
|
263
|
-
return
|
|
107
|
+
return LimiterHook(period, limit, get_key_with=lambda ctx: str(ctx.guild_id))
|
|
264
108
|
|
|
265
109
|
|
|
266
|
-
def channel_limiter(period: float, limit: int) ->
|
|
110
|
+
def channel_limiter(period: float, limit: int) -> LimiterHook[t.Any]:
|
|
267
111
|
"""Create a channel ratelimiter.
|
|
268
112
|
|
|
269
113
|
This ratelimiter is shared across all contexts in a channel.
|
|
@@ -274,11 +118,17 @@ def channel_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
|
274
118
|
The period, in seconds, after which the bucket resets.
|
|
275
119
|
limit : int
|
|
276
120
|
The amount of requests allowed in a bucket.
|
|
121
|
+
|
|
122
|
+
Usage
|
|
123
|
+
-----
|
|
124
|
+
```py
|
|
125
|
+
@arc.with_hook(arc.channel_limiter(5.0, 1))
|
|
126
|
+
```
|
|
277
127
|
"""
|
|
278
|
-
return
|
|
128
|
+
return LimiterHook(period, limit, get_key_with=lambda ctx: str(ctx.channel_id))
|
|
279
129
|
|
|
280
130
|
|
|
281
|
-
def user_limiter(period: float, limit: int) ->
|
|
131
|
+
def user_limiter(period: float, limit: int) -> LimiterHook[t.Any]:
|
|
282
132
|
"""Create a user ratelimiter.
|
|
283
133
|
|
|
284
134
|
This ratelimiter is shared across all contexts by a user.
|
|
@@ -289,11 +139,17 @@ def user_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
|
289
139
|
The period, in seconds, after which the bucket resets.
|
|
290
140
|
limit : int
|
|
291
141
|
The amount of requests allowed in a bucket.
|
|
142
|
+
|
|
143
|
+
Usage
|
|
144
|
+
-----
|
|
145
|
+
```py
|
|
146
|
+
@arc.with_hook(arc.user_limiter(5.0, 1))
|
|
147
|
+
```
|
|
292
148
|
"""
|
|
293
|
-
return
|
|
149
|
+
return LimiterHook(period, limit, get_key_with=lambda ctx: str(ctx.author.id))
|
|
294
150
|
|
|
295
151
|
|
|
296
|
-
def member_limiter(period: float, limit: int) ->
|
|
152
|
+
def member_limiter(period: float, limit: int) -> LimiterHook[t.Any]:
|
|
297
153
|
"""Create a member ratelimiter.
|
|
298
154
|
|
|
299
155
|
This ratelimiter is shared across all contexts by a member in a guild.
|
|
@@ -305,11 +161,17 @@ def member_limiter(period: float, limit: int) -> RateLimiter[t.Any]:
|
|
|
305
161
|
The period, in seconds, after which the bucket resets.
|
|
306
162
|
limit : int
|
|
307
163
|
The amount of requests allowed in a bucket.
|
|
164
|
+
|
|
165
|
+
Usage
|
|
166
|
+
-----
|
|
167
|
+
```py
|
|
168
|
+
@arc.with_hook(arc.member_limiter(5.0, 1))
|
|
169
|
+
```
|
|
308
170
|
"""
|
|
309
|
-
return
|
|
171
|
+
return LimiterHook(period, limit, get_key_with=lambda ctx: f"{ctx.author.id}:{ctx.guild_id}")
|
|
310
172
|
|
|
311
173
|
|
|
312
|
-
def custom_limiter(period: float, limit: int, get_key_with: t.Callable[[Context[t.Any]], str]) ->
|
|
174
|
+
def custom_limiter(period: float, limit: int, get_key_with: t.Callable[[Context[t.Any]], str]) -> LimiterHook[t.Any]:
|
|
313
175
|
"""Create a ratelimiter with a custom key extraction function.
|
|
314
176
|
|
|
315
177
|
Parameters
|
|
@@ -318,12 +180,18 @@ def custom_limiter(period: float, limit: int, get_key_with: t.Callable[[Context[
|
|
|
318
180
|
The period, in seconds, after which the bucket resets.
|
|
319
181
|
limit : int
|
|
320
182
|
The amount of requests allowed in a bucket.
|
|
321
|
-
get_key_with : Callable[[Context[t.Any]], str]
|
|
183
|
+
get_key_with : t.Callable[[Context[t.Any]], str]
|
|
322
184
|
A callable that returns a key for the ratelimiter bucket. This key is used to identify the bucket.
|
|
323
185
|
For instance, to create a ratelimiter that is shared across all contexts in a guild,
|
|
324
186
|
you would use `lambda ctx: str(ctx.guild_id)`.
|
|
187
|
+
|
|
188
|
+
Usage
|
|
189
|
+
-----
|
|
190
|
+
```py
|
|
191
|
+
@arc.with_hook(arc.custom_limiter(5.0, 1, lambda ctx: str(ctx.guild_id)))
|
|
192
|
+
```
|
|
325
193
|
"""
|
|
326
|
-
return
|
|
194
|
+
return LimiterHook(period, limit, get_key_with=get_key_with)
|
|
327
195
|
|
|
328
196
|
|
|
329
197
|
# MIT License
|
arc/utils/ratelimiter.py
ADDED
|
@@ -0,0 +1,243 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import sys
|
|
5
|
+
import time
|
|
6
|
+
import traceback
|
|
7
|
+
import typing as t
|
|
8
|
+
from collections import deque
|
|
9
|
+
|
|
10
|
+
import attr
|
|
11
|
+
|
|
12
|
+
from arc.abc.hookable import HookResult
|
|
13
|
+
|
|
14
|
+
KeyT = t.TypeVar("KeyT")
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@attr.define(slots=True, kw_only=True)
|
|
18
|
+
class _Bucket(t.Generic[KeyT]):
|
|
19
|
+
"""Handles the ratelimiting of a single item. (E.g. a single user or a channel)."""
|
|
20
|
+
|
|
21
|
+
key: str
|
|
22
|
+
"""The key of the bucket."""
|
|
23
|
+
|
|
24
|
+
reset_at: float
|
|
25
|
+
"""The time at which the bucket resets."""
|
|
26
|
+
|
|
27
|
+
limiter: RateLimiter[KeyT]
|
|
28
|
+
"""The limiter this bucket belongs to."""
|
|
29
|
+
|
|
30
|
+
_remaining: int = attr.field(alias="remaining")
|
|
31
|
+
"""The amount of requests remaining until the bucket is exhausted."""
|
|
32
|
+
|
|
33
|
+
_queue: deque[asyncio.Event] = attr.field(factory=deque, init=False)
|
|
34
|
+
"""A list of events to set as the iter task proceeds."""
|
|
35
|
+
|
|
36
|
+
_task: asyncio.Task[None] | None = attr.field(default=None, init=False)
|
|
37
|
+
"""The task that is currently iterating over the queue."""
|
|
38
|
+
|
|
39
|
+
@classmethod
|
|
40
|
+
def for_limiter(cls, key: str, limiter: RateLimiter[KeyT]) -> _Bucket[KeyT]:
|
|
41
|
+
"""Create a new bucket for a RateLimiter."""
|
|
42
|
+
return cls(key=key, limiter=limiter, reset_at=time.monotonic() + limiter.period, remaining=limiter.limit)
|
|
43
|
+
|
|
44
|
+
@property
|
|
45
|
+
def remaining(self) -> int:
|
|
46
|
+
"""The amount of requests remaining until the bucket is exhausted."""
|
|
47
|
+
if self.reset_at <= time.monotonic():
|
|
48
|
+
self.reset()
|
|
49
|
+
return self._remaining
|
|
50
|
+
|
|
51
|
+
@remaining.setter
|
|
52
|
+
def remaining(self, value: int) -> None:
|
|
53
|
+
self._remaining = value
|
|
54
|
+
|
|
55
|
+
@property
|
|
56
|
+
def is_exhausted(self) -> bool:
|
|
57
|
+
"""Return a boolean determining if the bucket is exhausted."""
|
|
58
|
+
return self.remaining <= 0 and self.reset_at > time.monotonic()
|
|
59
|
+
|
|
60
|
+
@property
|
|
61
|
+
def is_stale(self) -> bool:
|
|
62
|
+
"""Return a boolean determining if the bucket is stale.
|
|
63
|
+
If a bucket is stale, it is no longer in use and can be purged.
|
|
64
|
+
"""
|
|
65
|
+
return not self._queue and self.remaining == self.limiter.limit and (self._task is None or self._task.done())
|
|
66
|
+
|
|
67
|
+
def start_queue(self) -> None:
|
|
68
|
+
"""Start the queue of a bucket.
|
|
69
|
+
This will start setting events in the queue until the bucket is ratelimited.
|
|
70
|
+
"""
|
|
71
|
+
if self._task is None or self._task.done():
|
|
72
|
+
self._task = asyncio.create_task(self._iter_queue())
|
|
73
|
+
|
|
74
|
+
def reset(self) -> None:
|
|
75
|
+
"""Reset the bucket."""
|
|
76
|
+
self.reset_at = time.monotonic() + self.limiter.period
|
|
77
|
+
self._remaining = self.limiter.limit
|
|
78
|
+
|
|
79
|
+
async def _iter_queue(self) -> None:
|
|
80
|
+
"""Iterate over the queue and set events until exhausted."""
|
|
81
|
+
try:
|
|
82
|
+
while self._queue:
|
|
83
|
+
if self.remaining <= 0 and self.reset_at > time.monotonic():
|
|
84
|
+
# Sleep until ratelimit expires
|
|
85
|
+
await asyncio.sleep(self.reset_at - time.monotonic())
|
|
86
|
+
self.reset()
|
|
87
|
+
|
|
88
|
+
# Set events while not ratelimited
|
|
89
|
+
while self.remaining > 0 and self._queue:
|
|
90
|
+
self._queue.popleft().set()
|
|
91
|
+
self._remaining -= 1
|
|
92
|
+
|
|
93
|
+
except Exception as e:
|
|
94
|
+
print(f"Task Exception was never retrieved: {e}", file=sys.stderr)
|
|
95
|
+
print(traceback.format_exc(), file=sys.stderr)
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
class RateLimiter(t.Generic[KeyT]):
|
|
99
|
+
"""A fixed window ratelimiter implementation.
|
|
100
|
+
|
|
101
|
+
It is generic over the type of object it extracts the key from (`KeyT`), and must also be given a key extraction function.
|
|
102
|
+
It can be used independently for any purpose, or can be used as a hook via [`LimiterHook`][arc.utils.hooks.limiters.LimiterHook].
|
|
103
|
+
|
|
104
|
+
!!! tip
|
|
105
|
+
If you're looking to use this ratelimiter implementation as a hook, see [`LimiterHook`][arc.utils.hooks.limiters.LimiterHook]
|
|
106
|
+
along with all the other built-in limiter hooks.
|
|
107
|
+
|
|
108
|
+
Parameters
|
|
109
|
+
----------
|
|
110
|
+
period : float
|
|
111
|
+
The period, in seconds, after which the bucket resets.
|
|
112
|
+
limit : int
|
|
113
|
+
The amount of requests allowed in a bucket.
|
|
114
|
+
get_key_with : t.Callable[[Context[t.Any]], str]
|
|
115
|
+
A callable that returns a key for the ratelimiter bucket.
|
|
116
|
+
"""
|
|
117
|
+
|
|
118
|
+
__slots__ = ("period", "limit", "_buckets", "_get_key", "_gc_task")
|
|
119
|
+
|
|
120
|
+
def __init__(self, period: float, limit: int, *, get_key_with: t.Callable[[KeyT], str]) -> None:
|
|
121
|
+
self.period: float = period
|
|
122
|
+
self.limit: int = limit
|
|
123
|
+
self._buckets: t.Dict[str, _Bucket[KeyT]] = {}
|
|
124
|
+
self._get_key: t.Callable[[KeyT], str] = get_key_with
|
|
125
|
+
self._gc_task: asyncio.Task[None] | None = None
|
|
126
|
+
|
|
127
|
+
def get_key(self, ctx: KeyT) -> str:
|
|
128
|
+
"""Get key for ratelimiter bucket."""
|
|
129
|
+
return self._get_key(ctx)
|
|
130
|
+
|
|
131
|
+
def is_rate_limited(self, ctx: KeyT) -> bool:
|
|
132
|
+
"""Returns a boolean determining if the ratelimiter is ratelimited or not.
|
|
133
|
+
|
|
134
|
+
Parameters
|
|
135
|
+
----------
|
|
136
|
+
ctx : Context[t.Any]
|
|
137
|
+
The context to evaluate the ratelimit under.
|
|
138
|
+
|
|
139
|
+
Returns
|
|
140
|
+
-------
|
|
141
|
+
bool
|
|
142
|
+
A boolean determining if the ratelimiter is ratelimited or not.
|
|
143
|
+
"""
|
|
144
|
+
now = time.monotonic()
|
|
145
|
+
|
|
146
|
+
if data := self._buckets.get(self.get_key(ctx)):
|
|
147
|
+
if data.reset_at <= now:
|
|
148
|
+
return False
|
|
149
|
+
return data._remaining <= 0
|
|
150
|
+
return False
|
|
151
|
+
|
|
152
|
+
def _start_gc(self) -> None:
|
|
153
|
+
"""Start the garbage collector task if one is not running."""
|
|
154
|
+
if self._gc_task is None or self._gc_task.done():
|
|
155
|
+
self._gc_task = asyncio.create_task(self._gc())
|
|
156
|
+
|
|
157
|
+
async def _gc(self) -> None:
|
|
158
|
+
"""Purge stale buckets."""
|
|
159
|
+
while self._buckets:
|
|
160
|
+
await asyncio.sleep(self.period + 1.0)
|
|
161
|
+
for bucket in list(self._buckets.values()):
|
|
162
|
+
if bucket.is_stale:
|
|
163
|
+
del self._buckets[bucket.key]
|
|
164
|
+
|
|
165
|
+
async def acquire(self, ctx: KeyT, *, wait: bool = True) -> None:
|
|
166
|
+
"""Acquire a bucket, block execution if ratelimited and wait is True.
|
|
167
|
+
|
|
168
|
+
Parameters
|
|
169
|
+
----------
|
|
170
|
+
ctx : Context[t.Any]
|
|
171
|
+
The context to evaluate the ratelimit under.
|
|
172
|
+
wait : bool
|
|
173
|
+
Determines if this call should block in
|
|
174
|
+
case of hitting a ratelimit.
|
|
175
|
+
|
|
176
|
+
Raises
|
|
177
|
+
------
|
|
178
|
+
RateLimiterExhaustedError
|
|
179
|
+
If the ratelimiter is ratelimited and wait is False.
|
|
180
|
+
"""
|
|
181
|
+
event = asyncio.Event()
|
|
182
|
+
|
|
183
|
+
key = self.get_key(ctx)
|
|
184
|
+
# Get existing or insert new bucket
|
|
185
|
+
bucket = self._buckets.setdefault(key, _Bucket.for_limiter(key, self))
|
|
186
|
+
|
|
187
|
+
if bucket.is_exhausted and not wait:
|
|
188
|
+
raise RateLimiterExhaustedError(
|
|
189
|
+
self,
|
|
190
|
+
bucket.reset_at - time.monotonic(),
|
|
191
|
+
f"Ratelimited for {bucket.reset_at - time.monotonic()} seconds.",
|
|
192
|
+
)
|
|
193
|
+
|
|
194
|
+
bucket._queue.append(event)
|
|
195
|
+
bucket.start_queue()
|
|
196
|
+
self._start_gc()
|
|
197
|
+
|
|
198
|
+
if wait:
|
|
199
|
+
await event.wait()
|
|
200
|
+
|
|
201
|
+
async def __call__(self, ctx: KeyT) -> HookResult:
|
|
202
|
+
"""Acquire a ratelimit, fail if ratelimited.
|
|
203
|
+
|
|
204
|
+
Parameters
|
|
205
|
+
----------
|
|
206
|
+
ctx : Context[t.Any]
|
|
207
|
+
The context to evaluate the ratelimit under.
|
|
208
|
+
|
|
209
|
+
Returns
|
|
210
|
+
-------
|
|
211
|
+
HookResult
|
|
212
|
+
A hook result to conform to the limiter protocol.
|
|
213
|
+
|
|
214
|
+
Raises
|
|
215
|
+
------
|
|
216
|
+
UnderCooldownError
|
|
217
|
+
If the ratelimiter is ratelimited.
|
|
218
|
+
"""
|
|
219
|
+
await self.acquire(ctx, wait=False)
|
|
220
|
+
return HookResult()
|
|
221
|
+
|
|
222
|
+
def reset(self, ctx: KeyT) -> None:
|
|
223
|
+
"""Reset the ratelimit for a given context."""
|
|
224
|
+
if bucket := self._buckets.get(self.get_key(ctx)):
|
|
225
|
+
bucket.reset()
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
class RateLimiterExhaustedError(Exception):
|
|
229
|
+
"""Raised when a [RateLimiter][arc.utils.ratelimiter.RateLimiter] is acquired while it is exhausted, and the
|
|
230
|
+
`wait` parameter is set to `False`.
|
|
231
|
+
|
|
232
|
+
Attributes
|
|
233
|
+
----------
|
|
234
|
+
retry_after : float
|
|
235
|
+
The amount of time in seconds until the command is off cooldown.
|
|
236
|
+
limiter : arc.abc.limiter.LimiterProto
|
|
237
|
+
The limiter that was rate limited.
|
|
238
|
+
"""
|
|
239
|
+
|
|
240
|
+
def __init__(self, limiter: RateLimiter[t.Any], retry_after: float, *args: t.Any) -> None:
|
|
241
|
+
self.retry_after = retry_after
|
|
242
|
+
self.limiter = limiter
|
|
243
|
+
super().__init__(*args)
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: hikari-arc
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.6.0
|
|
4
4
|
Summary: A command handler for hikari with a focus on type-safety and correctness.
|
|
5
5
|
Home-page: https://github.com/hypergonial/hikari-arc
|
|
6
6
|
Author: hypergonial
|
|
@@ -29,12 +29,17 @@ Requires-Dist: attrs >=23.1
|
|
|
29
29
|
Requires-Dist: colorama ; sys_platform=="win32"
|
|
30
30
|
Provides-Extra: dev
|
|
31
31
|
Requires-Dist: ruff ==0.1.11 ; extra == 'dev'
|
|
32
|
-
Requires-Dist: pyright ==1.1.
|
|
32
|
+
Requires-Dist: pyright ==1.1.345 ; extra == 'dev'
|
|
33
33
|
Requires-Dist: nox ==2023.4.22 ; extra == 'dev'
|
|
34
|
+
Requires-Dist: typing-extensions ==4.9.0 ; extra == 'dev'
|
|
35
|
+
Requires-Dist: pytest ==7.4.4 ; extra == 'dev'
|
|
36
|
+
Requires-Dist: pytest-asyncio ==0.23.3 ; extra == 'dev'
|
|
34
37
|
Provides-Extra: docs
|
|
35
38
|
Requires-Dist: mkdocs-material[imaging] ~=9.5.3 ; extra == 'docs'
|
|
36
39
|
Requires-Dist: mkdocs ~=1.5.3 ; extra == 'docs'
|
|
37
|
-
Requires-Dist: mkdocstrings-python ~=1.
|
|
40
|
+
Requires-Dist: mkdocstrings-python ~=1.8.0 ; extra == 'docs'
|
|
41
|
+
Requires-Dist: black ~=23.12.1 ; extra == 'docs'
|
|
42
|
+
Requires-Dist: griffe-inherited-docstrings ~=1.0.0 ; extra == 'docs'
|
|
38
43
|
Provides-Extra: rest
|
|
39
44
|
Requires-Dist: hikari[server] >=2.0.0.dev122 ; extra == 'rest'
|
|
40
45
|
|
|
@@ -105,9 +110,9 @@ See [Contributing](./CONTRIBUTING.md)
|
|
|
105
110
|
`arc` is in large part a combination of all the parts I like in other command handlers, with my own spin on it. The following projects have inspired me and aided me greatly in the design of this library:
|
|
106
111
|
|
|
107
112
|
- [`hikari-lightbulb`](https://github.com/tandemdude/hikari-lightbulb) - The library initially started as a reimagination of lightbulb, it inherits a similar project structure and terminology.
|
|
108
|
-
- [`Tanjun`](https://github.com/FasterSpeeding/Tanjun) - For the idea of using `typing.Annotated` and dependency injection in a command handler. `arc` also uses the same dependency injection library, [`Alluka`](https://github.com/FasterSpeeding/Alluka), under the hood.
|
|
109
|
-
- [`hikari-crescent`](https://github.com/hikari-crescent/hikari-crescent) The design of hooks is largely inspired by `crescent`.
|
|
110
|
-
- [`FastAPI`](https://github.com/tiangolo/fastapi) - Some design ideas and most of the documentation configuration derives from `FastAPI`.
|
|
113
|
+
- [`Tanjun`](https://github.com/FasterSpeeding/Tanjun) - For the idea of using `typing.Annotated` and [dependency injection](https://arc.hypergonial.com/guides/dependency_injection/) in a command handler. `arc` also uses the same dependency injection library, [`Alluka`](https://github.com/FasterSpeeding/Alluka), under the hood.
|
|
114
|
+
- [`hikari-crescent`](https://github.com/hikari-crescent/hikari-crescent) The design of [hooks](https://arc.hypergonial.com/guides/hooks/) is largely inspired by `crescent`.
|
|
115
|
+
- [`FastAPI`](https://github.com/tiangolo/fastapi) - Some design ideas and most of the [documentation](https://arc.hypergonial.com/) [configuration](https://github.com/hypergonial/hikari-arc/blob/main/mkdocs.yml) derives from `FastAPI`.
|
|
111
116
|
|
|
112
117
|
## Links
|
|
113
118
|
|