python-wrap 0.2.1__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.
- python_wrap/__init__.py +272 -0
- python_wrap/_core.py +212 -0
- python_wrap/aio.py +620 -0
- python_wrap/batching.py +556 -0
- python_wrap/bulkhead.py +431 -0
- python_wrap/caching.py +1613 -0
- python_wrap/circuit.py +553 -0
- python_wrap/concurrency.py +1886 -0
- python_wrap/config.py +217 -0
- python_wrap/deadline.py +168 -0
- python_wrap/decorators.py +7 -0
- python_wrap/devtools.py +385 -0
- python_wrap/errors.py +86 -0
- python_wrap/events.py +161 -0
- python_wrap/hedging.py +192 -0
- python_wrap/logs.py +767 -0
- python_wrap/metrics.py +242 -0
- python_wrap/periodic.py +491 -0
- python_wrap/policy.py +263 -0
- python_wrap/py.typed +0 -0
- python_wrap/retry.py +1324 -0
- python_wrap/state.py +500 -0
- python_wrap/testing.py +662 -0
- python_wrap/timing.py +535 -0
- python_wrap/validation.py +662 -0
- python_wrap-0.2.1.dist-info/METADATA +485 -0
- python_wrap-0.2.1.dist-info/RECORD +29 -0
- python_wrap-0.2.1.dist-info/WHEEL +4 -0
- python_wrap-0.2.1.dist-info/licenses/LICENSE +21 -0
|
@@ -0,0 +1,1886 @@
|
|
|
1
|
+
"""Concurrency and flow-control decorators.
|
|
2
|
+
|
|
3
|
+
* `timeout` — bound how long a call (or a `with` / `async with` block) may run.
|
|
4
|
+
* `rate_limit` — at most *N* calls per period, rejecting or waiting.
|
|
5
|
+
* `RateLimiter` — a rate limit shared by several functions and blocks.
|
|
6
|
+
* `run_in_thread` — start a call in a new thread and get a ``Future`` back.
|
|
7
|
+
* `throttle` — run at most once per interval; calls in between are skipped.
|
|
8
|
+
* `debounce` — run once activity has stopped for a while.
|
|
9
|
+
* `singleflight` — concurrent identical calls share one execution.
|
|
10
|
+
|
|
11
|
+
Every decorator works on plain functions, ``async def`` functions and methods. State
|
|
12
|
+
kept by a decorator (rate-limit buckets, throttle windows, pending debounced calls) belongs
|
|
13
|
+
to the decorated function, so on a method it is shared by every instance unless a ``key``
|
|
14
|
+
argument says otherwise.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import asyncio
|
|
20
|
+
import concurrent.futures
|
|
21
|
+
import contextlib
|
|
22
|
+
import contextvars
|
|
23
|
+
import functools
|
|
24
|
+
import os
|
|
25
|
+
import queue
|
|
26
|
+
import signal
|
|
27
|
+
import sys
|
|
28
|
+
import threading
|
|
29
|
+
import time
|
|
30
|
+
import weakref
|
|
31
|
+
from collections import deque
|
|
32
|
+
from collections.abc import Awaitable, Callable, Coroutine, Hashable, Iterable
|
|
33
|
+
from types import FrameType, TracebackType
|
|
34
|
+
from typing import (
|
|
35
|
+
Any,
|
|
36
|
+
Concatenate,
|
|
37
|
+
Final,
|
|
38
|
+
Generic,
|
|
39
|
+
NoReturn,
|
|
40
|
+
ParamSpec,
|
|
41
|
+
Protocol,
|
|
42
|
+
TypeVar,
|
|
43
|
+
cast,
|
|
44
|
+
overload,
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
from python_wrap import _core
|
|
48
|
+
from python_wrap import deadline as _deadline
|
|
49
|
+
from python_wrap._core import F, P, R
|
|
50
|
+
from python_wrap.errors import RateLimitExceeded, TimeoutExpired
|
|
51
|
+
|
|
52
|
+
__all__ = [
|
|
53
|
+
"AsyncDebounced",
|
|
54
|
+
"Debounced",
|
|
55
|
+
"Limit",
|
|
56
|
+
"RateLimiter",
|
|
57
|
+
"Timeout",
|
|
58
|
+
"debounce",
|
|
59
|
+
"limiter_of",
|
|
60
|
+
"rate_limit",
|
|
61
|
+
"run_in_thread",
|
|
62
|
+
"singleflight",
|
|
63
|
+
"throttle",
|
|
64
|
+
"timeout",
|
|
65
|
+
]
|
|
66
|
+
|
|
67
|
+
T = TypeVar("T")
|
|
68
|
+
R_co = TypeVar("R_co", covariant=True)
|
|
69
|
+
P2 = ParamSpec("P2")
|
|
70
|
+
|
|
71
|
+
# ---------------------------------------------------------------------------------------
|
|
72
|
+
# timeout
|
|
73
|
+
# ---------------------------------------------------------------------------------------
|
|
74
|
+
|
|
75
|
+
_HAS_SETITIMER = hasattr(signal, "setitimer") and hasattr(signal, "SIGALRM")
|
|
76
|
+
_BLOCK = "<block>" # the "function" name used in messages and events for `with timeout()`
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def _signals_usable() -> bool:
|
|
80
|
+
return _HAS_SETITIMER and threading.current_thread() is threading.main_thread()
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _timeout_message(label: str, seconds: float) -> str:
|
|
84
|
+
return f"{label} did not finish within {seconds:g}s"
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def _expired(label: str, seconds: float) -> TimeoutExpired:
|
|
88
|
+
"""Build the error for an expiry and emit ``timeout.expired``."""
|
|
89
|
+
_core.emit("timeout.expired", label, seconds=seconds)
|
|
90
|
+
return TimeoutExpired(_timeout_message(label, seconds), seconds=seconds)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
class _Alarm:
|
|
94
|
+
"""One SIGALRM/``setitimer`` deadline, cooperating with enclosing ones.
|
|
95
|
+
|
|
96
|
+
Enter it in the main thread only. The previous handler and any armed timer are
|
|
97
|
+
restored on exit; an enclosing timer whose deadline comes first still fires.
|
|
98
|
+
Entering costs two system calls (install the handler, arm the timer — arming also
|
|
99
|
+
reports the enclosing timer, so no separate ``getitimer``) and so does exiting.
|
|
100
|
+
"""
|
|
101
|
+
|
|
102
|
+
def __init__(self, label: str, seconds: float) -> None:
|
|
103
|
+
self.label = label
|
|
104
|
+
self.seconds = seconds
|
|
105
|
+
self.error: TimeoutExpired | None = None # built only when the alarm fires
|
|
106
|
+
self._outer_fired = False
|
|
107
|
+
# The enclosing timer's delay; None until known (our timer is not armed yet).
|
|
108
|
+
self._previous_delay: float | None = None
|
|
109
|
+
self._previous_interval = 0.0
|
|
110
|
+
|
|
111
|
+
def __enter__(self) -> None:
|
|
112
|
+
# The itimer counts real time, so the deadline arithmetic does too (a fake clock
|
|
113
|
+
# installed by tests must not move it).
|
|
114
|
+
self._start = time.monotonic()
|
|
115
|
+
deadline = self._start + self.seconds
|
|
116
|
+
self._previous_handler = previous_handler = signal.getsignal(signal.SIGALRM)
|
|
117
|
+
|
|
118
|
+
def handler(signum: int, frame: FrameType | None) -> None:
|
|
119
|
+
previous_delay = self._previous_delay
|
|
120
|
+
if previous_delay is not None and (
|
|
121
|
+
not previous_delay or time.monotonic() >= deadline - 1e-3
|
|
122
|
+
):
|
|
123
|
+
self.error = TimeoutExpired(
|
|
124
|
+
_timeout_message(self.label, self.seconds), seconds=self.seconds
|
|
125
|
+
)
|
|
126
|
+
raise self.error
|
|
127
|
+
# The alarm belongs to an enclosing timeout whose deadline comes first (or
|
|
128
|
+
# fired before ours was even armed).
|
|
129
|
+
self._outer_fired = True
|
|
130
|
+
signal.signal(signal.SIGALRM, previous_handler)
|
|
131
|
+
if callable(previous_handler):
|
|
132
|
+
previous_handler(signum, frame)
|
|
133
|
+
|
|
134
|
+
signal.signal(signal.SIGALRM, handler)
|
|
135
|
+
previous_delay, self._previous_interval = signal.setitimer(signal.ITIMER_REAL, self.seconds)
|
|
136
|
+
if previous_delay and previous_delay < self.seconds:
|
|
137
|
+
# An enclosing timer is due first: keep its deadline.
|
|
138
|
+
signal.setitimer(signal.ITIMER_REAL, previous_delay)
|
|
139
|
+
self._previous_delay = previous_delay
|
|
140
|
+
|
|
141
|
+
def __exit__(
|
|
142
|
+
self,
|
|
143
|
+
exc_type: type[BaseException] | None,
|
|
144
|
+
exc: BaseException | None,
|
|
145
|
+
tb: TracebackType | None,
|
|
146
|
+
) -> None:
|
|
147
|
+
signal.setitimer(signal.ITIMER_REAL, 0)
|
|
148
|
+
if not self._outer_fired:
|
|
149
|
+
signal.signal(signal.SIGALRM, self._previous_handler)
|
|
150
|
+
if self._previous_delay:
|
|
151
|
+
remaining = self._previous_delay - (time.monotonic() - self._start)
|
|
152
|
+
# Re-arm the enclosing timer; if its deadline already passed, fire ASAP.
|
|
153
|
+
signal.setitimer(signal.ITIMER_REAL, max(remaining, 1e-6), self._previous_interval)
|
|
154
|
+
if exc is not None and exc is self.error:
|
|
155
|
+
_core.emit("timeout.expired", self.label, seconds=self.seconds)
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
# -- thread fallback: a small pool of idle daemon workers ---------------------------------
|
|
159
|
+
|
|
160
|
+
_WORKER_IDLE_SECONDS = 5.0 # an idle worker exits after this long (real time)
|
|
161
|
+
_MAX_IDLE_WORKERS = 8 # idle workers kept for reuse; extra ones exit when done
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
class _Job:
|
|
165
|
+
"""One call handed to a worker; ``done`` is released when it has finished."""
|
|
166
|
+
|
|
167
|
+
__slots__ = (
|
|
168
|
+
"abandoned",
|
|
169
|
+
"args",
|
|
170
|
+
"context",
|
|
171
|
+
"done",
|
|
172
|
+
"error",
|
|
173
|
+
"finished",
|
|
174
|
+
"func",
|
|
175
|
+
"kwargs",
|
|
176
|
+
"result",
|
|
177
|
+
)
|
|
178
|
+
|
|
179
|
+
def __init__(self, func: Callable[..., Any], args: Any, kwargs: Any) -> None:
|
|
180
|
+
self.func = func
|
|
181
|
+
self.args = args
|
|
182
|
+
self.kwargs = kwargs
|
|
183
|
+
self.context = contextvars.copy_context() # per call, as with a fresh thread
|
|
184
|
+
self.done = threading.Lock()
|
|
185
|
+
self.done.acquire()
|
|
186
|
+
self.result: Any = None
|
|
187
|
+
self.error: BaseException | None = None
|
|
188
|
+
self.abandoned = False # the caller timed out and stopped waiting
|
|
189
|
+
self.finished = 0.0 # `time.monotonic()` when the call returned (real time)
|
|
190
|
+
|
|
191
|
+
def run(self) -> None:
|
|
192
|
+
try:
|
|
193
|
+
self.result = self.context.run(self.func, *self.args, **self.kwargs)
|
|
194
|
+
except BaseException as exc:
|
|
195
|
+
self.error = exc
|
|
196
|
+
self.func = self.args = self.kwargs = None # type: ignore[assignment]
|
|
197
|
+
self.finished = time.monotonic()
|
|
198
|
+
self.done.release()
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
class _WorkerPool:
|
|
202
|
+
"""Reusable daemon threads for the `timeout` thread fallback.
|
|
203
|
+
|
|
204
|
+
A finished worker goes back to the idle list (up to ``_MAX_IDLE_WORKERS``); a worker
|
|
205
|
+
whose call timed out is abandoned — it exits when that call eventually returns and
|
|
206
|
+
is never reused. An idle worker exits after ``_WORKER_IDLE_SECONDS`` without work.
|
|
207
|
+
"""
|
|
208
|
+
|
|
209
|
+
def __init__(self) -> None:
|
|
210
|
+
self.lock = threading.Lock()
|
|
211
|
+
self.idle: list[queue.SimpleQueue[_Job]] = [] # inboxes of idle workers (LIFO)
|
|
212
|
+
self.count = 0
|
|
213
|
+
|
|
214
|
+
def submit(self, job: _Job) -> None:
|
|
215
|
+
with self.lock:
|
|
216
|
+
inbox = self.idle.pop() if self.idle else None
|
|
217
|
+
if inbox is None:
|
|
218
|
+
self.count += 1
|
|
219
|
+
number = self.count
|
|
220
|
+
if inbox is not None:
|
|
221
|
+
inbox.put(job)
|
|
222
|
+
return
|
|
223
|
+
inbox = queue.SimpleQueue()
|
|
224
|
+
inbox.put(job)
|
|
225
|
+
threading.Thread(
|
|
226
|
+
target=self._work, args=(inbox,), name=f"timeout-worker-{number}", daemon=True
|
|
227
|
+
).start()
|
|
228
|
+
|
|
229
|
+
def _work(self, inbox: queue.SimpleQueue[_Job]) -> None:
|
|
230
|
+
while True:
|
|
231
|
+
try:
|
|
232
|
+
job = inbox.get(timeout=_WORKER_IDLE_SECONDS)
|
|
233
|
+
except queue.Empty:
|
|
234
|
+
with self.lock:
|
|
235
|
+
if inbox in self.idle:
|
|
236
|
+
self.idle.remove(inbox)
|
|
237
|
+
return
|
|
238
|
+
continue # claimed just as we gave up waiting: the job is on its way
|
|
239
|
+
job.run()
|
|
240
|
+
with self.lock:
|
|
241
|
+
if job.abandoned or len(self.idle) >= _MAX_IDLE_WORKERS:
|
|
242
|
+
return
|
|
243
|
+
self.idle.append(inbox)
|
|
244
|
+
del job
|
|
245
|
+
|
|
246
|
+
def abandon(self, job: _Job) -> None:
|
|
247
|
+
with self.lock:
|
|
248
|
+
job.abandoned = True
|
|
249
|
+
|
|
250
|
+
def reset(self) -> None:
|
|
251
|
+
"""After `os.fork`: the parent's workers do not exist in the child."""
|
|
252
|
+
self.lock = threading.Lock()
|
|
253
|
+
self.idle = []
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
_workers = _WorkerPool()
|
|
257
|
+
if hasattr(os, "register_at_fork"): # pragma: no branch - absent on Windows only
|
|
258
|
+
os.register_at_fork(after_in_child=_workers.reset)
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
def _call_in_worker(
|
|
262
|
+
func: Callable[P, R], label: str, seconds: float, *args: P.args, **kwargs: P.kwargs
|
|
263
|
+
) -> R:
|
|
264
|
+
"""Run ``func`` in a pooled daemon thread and stop waiting for it after ``seconds``."""
|
|
265
|
+
job = _Job(func, args, kwargs)
|
|
266
|
+
# Real time, like the wait itself: a fake clock advanced by sleeps inside the call
|
|
267
|
+
# (e.g. a retry) must not turn a quick call into a late one.
|
|
268
|
+
start = time.monotonic()
|
|
269
|
+
_workers.submit(job)
|
|
270
|
+
if not job.done.acquire(timeout=seconds):
|
|
271
|
+
_workers.abandon(job)
|
|
272
|
+
raise _expired(label, seconds) from None
|
|
273
|
+
if job.finished - start > seconds:
|
|
274
|
+
# The wait overslept (seen with free-threaded builds) and the call finished
|
|
275
|
+
# late: it still ran over its limit.
|
|
276
|
+
raise _expired(label, seconds) from None
|
|
277
|
+
error = job.error
|
|
278
|
+
if error is not None:
|
|
279
|
+
job.error = None
|
|
280
|
+
try:
|
|
281
|
+
raise error
|
|
282
|
+
finally:
|
|
283
|
+
del error
|
|
284
|
+
return cast(R, job.result)
|
|
285
|
+
|
|
286
|
+
|
|
287
|
+
# Message of the cancellations a timer sends itself. Python 3.10 has no
|
|
288
|
+
# ``Task.cancelling()`` counter, so there the message is how a timer tells its own
|
|
289
|
+
# cancellation apart from one sent by someone else (best effort: two cancellations that
|
|
290
|
+
# arrive before the task resumes are delivered as one, as in the 3.10 stdlib).
|
|
291
|
+
_TIMER_CANCEL = "python_wrap.timeout expired"
|
|
292
|
+
|
|
293
|
+
|
|
294
|
+
def _is_timer_cancel(exc: BaseException | None) -> bool:
|
|
295
|
+
"""True if ``exc`` is the CancelledError a timer sent (3.10 path only)."""
|
|
296
|
+
return isinstance(exc, asyncio.CancelledError) and exc.args[:1] == (_TIMER_CANCEL,)
|
|
297
|
+
|
|
298
|
+
|
|
299
|
+
if sys.version_info >= (3, 11):
|
|
300
|
+
_async_timer = asyncio.timeout
|
|
301
|
+
|
|
302
|
+
else: # pragma: no cover - exercised by the Python 3.10 CI job
|
|
303
|
+
|
|
304
|
+
class _CancelTimer:
|
|
305
|
+
"""`asyncio.timeout` stand-in for 3.10: cancel the task at the deadline.
|
|
306
|
+
|
|
307
|
+
Like `asyncio.timeout`, the cancellation that it caused is turned into a
|
|
308
|
+
``TimeoutError`` on exit; any other cancellation propagates unchanged.
|
|
309
|
+
"""
|
|
310
|
+
|
|
311
|
+
def __init__(self, seconds: float) -> None:
|
|
312
|
+
self._seconds = seconds
|
|
313
|
+
self._fired = False
|
|
314
|
+
|
|
315
|
+
def expired(self) -> bool:
|
|
316
|
+
return self._fired
|
|
317
|
+
|
|
318
|
+
async def __aenter__(self) -> _CancelTimer:
|
|
319
|
+
task = asyncio.current_task()
|
|
320
|
+
if task is None:
|
|
321
|
+
raise RuntimeError("timeout() must be used inside a task")
|
|
322
|
+
|
|
323
|
+
def fire() -> None:
|
|
324
|
+
self._fired = True
|
|
325
|
+
task.cancel(_TIMER_CANCEL)
|
|
326
|
+
|
|
327
|
+
self._handle = asyncio.get_running_loop().call_later(self._seconds, fire)
|
|
328
|
+
return self
|
|
329
|
+
|
|
330
|
+
async def __aexit__(
|
|
331
|
+
self,
|
|
332
|
+
exc_type: type[BaseException] | None,
|
|
333
|
+
exc: BaseException | None,
|
|
334
|
+
tb: TracebackType | None,
|
|
335
|
+
) -> None:
|
|
336
|
+
self._handle.cancel()
|
|
337
|
+
if self._fired and _is_timer_cancel(exc):
|
|
338
|
+
raise TimeoutError from exc
|
|
339
|
+
|
|
340
|
+
_async_timer = _CancelTimer
|
|
341
|
+
|
|
342
|
+
|
|
343
|
+
class _GraceTimer:
|
|
344
|
+
"""Cancel the current task at the deadline, and again after ``grace`` seconds.
|
|
345
|
+
|
|
346
|
+
The first cancellation lets the block run its ``finally`` / ``except CancelledError``
|
|
347
|
+
cleanup; if the cleanup is still awaiting when the grace period ends, a second
|
|
348
|
+
cancellation interrupts it (with a warning). Only cancellations caused by this timer
|
|
349
|
+
are turned into ``TimeoutError``.
|
|
350
|
+
"""
|
|
351
|
+
|
|
352
|
+
def __init__(self, seconds: float, grace: float, label: str) -> None:
|
|
353
|
+
self._seconds = seconds
|
|
354
|
+
self._grace = grace
|
|
355
|
+
self._label = label
|
|
356
|
+
self._fired = False
|
|
357
|
+
self._cancels = 0
|
|
358
|
+
self._cancelling = 0
|
|
359
|
+
self._handles: list[asyncio.TimerHandle] = []
|
|
360
|
+
|
|
361
|
+
def expired(self) -> bool:
|
|
362
|
+
return self._fired
|
|
363
|
+
|
|
364
|
+
async def __aenter__(self) -> _GraceTimer:
|
|
365
|
+
task = asyncio.current_task()
|
|
366
|
+
if task is None: # pragma: no cover - asyncio always runs coroutines in a task
|
|
367
|
+
raise RuntimeError("timeout() must be used inside a task")
|
|
368
|
+
loop = asyncio.get_running_loop()
|
|
369
|
+
|
|
370
|
+
def second() -> None:
|
|
371
|
+
self._cancels += 1
|
|
372
|
+
_core.get_logger(None).warning(
|
|
373
|
+
"%s: cleanup still running %ss after the timeout, cancelling it again",
|
|
374
|
+
self._label,
|
|
375
|
+
self._grace,
|
|
376
|
+
)
|
|
377
|
+
task.cancel(_TIMER_CANCEL)
|
|
378
|
+
|
|
379
|
+
def fire() -> None:
|
|
380
|
+
self._fired = True
|
|
381
|
+
self._cancels += 1
|
|
382
|
+
task.cancel(_TIMER_CANCEL)
|
|
383
|
+
self._handles.append(loop.call_later(self._grace, second))
|
|
384
|
+
|
|
385
|
+
self._task = task
|
|
386
|
+
# Like asyncio.timeout: remember the cancel requests already pending on entry, so
|
|
387
|
+
# on exit only the timer's own ones are converted (3.11+).
|
|
388
|
+
cancelling = getattr(task, "cancelling", None)
|
|
389
|
+
self._cancelling = cancelling() if cancelling is not None else 0
|
|
390
|
+
self._handles.append(loop.call_later(self._seconds, fire))
|
|
391
|
+
return self
|
|
392
|
+
|
|
393
|
+
async def __aexit__(
|
|
394
|
+
self,
|
|
395
|
+
exc_type: type[BaseException] | None,
|
|
396
|
+
exc: BaseException | None,
|
|
397
|
+
tb: TracebackType | None,
|
|
398
|
+
) -> None:
|
|
399
|
+
for handle in self._handles:
|
|
400
|
+
handle.cancel()
|
|
401
|
+
if not self._fired or exc_type is not asyncio.CancelledError:
|
|
402
|
+
return
|
|
403
|
+
uncancel = getattr(self._task, "uncancel", None)
|
|
404
|
+
if uncancel is not None: # 3.11+: withdraw our own requests, then count the rest
|
|
405
|
+
remaining: int = getattr(self._task, "cancelling")() # noqa: B009 - 3.11+ only
|
|
406
|
+
for _ in range(self._cancels):
|
|
407
|
+
if remaining <= self._cancelling:
|
|
408
|
+
break
|
|
409
|
+
remaining = uncancel()
|
|
410
|
+
if remaining <= self._cancelling:
|
|
411
|
+
raise TimeoutError from exc
|
|
412
|
+
elif _is_timer_cancel(exc): # pragma: no cover - exercised by the 3.10 CI job
|
|
413
|
+
raise TimeoutError from exc
|
|
414
|
+
|
|
415
|
+
|
|
416
|
+
async def _await_with_grace(awaitable: Awaitable[T], label: str, seconds: float, grace: float) -> T:
|
|
417
|
+
"""Run ``awaitable`` as a child task; on expiry cancel it and wait ``grace`` for cleanup."""
|
|
418
|
+
task = asyncio.ensure_future(awaitable)
|
|
419
|
+
try:
|
|
420
|
+
await asyncio.wait({task}, timeout=seconds)
|
|
421
|
+
except BaseException:
|
|
422
|
+
# Cancelled from outside: cancel the child too, give its cleanup the same grace,
|
|
423
|
+
# then re-raise the caller's cancellation. Never leave its outcome unretrieved.
|
|
424
|
+
task.cancel()
|
|
425
|
+
try:
|
|
426
|
+
await asyncio.wait({task}, timeout=grace)
|
|
427
|
+
finally:
|
|
428
|
+
_consume_when_done(task)
|
|
429
|
+
raise
|
|
430
|
+
if task.done():
|
|
431
|
+
return task.result()
|
|
432
|
+
task.cancel()
|
|
433
|
+
try:
|
|
434
|
+
await asyncio.wait({task}, timeout=grace)
|
|
435
|
+
except BaseException:
|
|
436
|
+
_consume_when_done(task) # cancelled from outside during the grace period
|
|
437
|
+
raise
|
|
438
|
+
if not task.done():
|
|
439
|
+
_core.get_logger(None).warning(
|
|
440
|
+
"%s: still running %ss after the timeout; abandoning it", label, grace
|
|
441
|
+
)
|
|
442
|
+
_consume_when_done(task)
|
|
443
|
+
raise _expired(label, seconds)
|
|
444
|
+
|
|
445
|
+
|
|
446
|
+
def _consume_when_done(task: asyncio.Future[Any]) -> None:
|
|
447
|
+
"""Retrieve ``task``'s outcome now if it is done, otherwise once it is."""
|
|
448
|
+
if task.done():
|
|
449
|
+
_consume_result(task)
|
|
450
|
+
else:
|
|
451
|
+
task.add_done_callback(_consume_result)
|
|
452
|
+
|
|
453
|
+
|
|
454
|
+
def _consume_result(task: asyncio.Future[Any]) -> None:
|
|
455
|
+
"""Mark an abandoned task's outcome as retrieved so asyncio does not log it."""
|
|
456
|
+
if not task.cancelled():
|
|
457
|
+
task.exception()
|
|
458
|
+
|
|
459
|
+
|
|
460
|
+
async def _await_with_deadline(awaitable: Awaitable[T], label: str, seconds: float) -> T:
|
|
461
|
+
"""Await under a timer; convert only *our* expiry to TimeoutExpired."""
|
|
462
|
+
timer = _async_timer(seconds)
|
|
463
|
+
try:
|
|
464
|
+
async with timer:
|
|
465
|
+
return await awaitable
|
|
466
|
+
except TimeoutError as exc:
|
|
467
|
+
if not timer.expired():
|
|
468
|
+
raise # raised by the wrapped function itself: propagate unchanged
|
|
469
|
+
raise _expired(label, seconds) from exc
|
|
470
|
+
|
|
471
|
+
|
|
472
|
+
class Timeout:
|
|
473
|
+
"""A time limit usable as a decorator, `with` block or `async with` block.
|
|
474
|
+
|
|
475
|
+
Created by `timeout`; see there for the strategies and caveats. One instance can
|
|
476
|
+
decorate several functions and be entered several times (also concurrently from
|
|
477
|
+
different tasks); every use gets its own timer.
|
|
478
|
+
|
|
479
|
+
Attributes:
|
|
480
|
+
seconds: The configured limit.
|
|
481
|
+
use_signals: The configured sync strategy (`None` means automatic).
|
|
482
|
+
cleanup_grace: Seconds async code may spend cleaning up after the limit (`None`: no
|
|
483
|
+
separate bound).
|
|
484
|
+
"""
|
|
485
|
+
|
|
486
|
+
def __init__(
|
|
487
|
+
self,
|
|
488
|
+
seconds: float,
|
|
489
|
+
*,
|
|
490
|
+
use_signals: bool | None = None,
|
|
491
|
+
cleanup_grace: float | None = None,
|
|
492
|
+
) -> None:
|
|
493
|
+
if seconds <= 0:
|
|
494
|
+
raise ValueError(f"timeout seconds must be positive, got {seconds!r}")
|
|
495
|
+
if cleanup_grace is not None and cleanup_grace < 0:
|
|
496
|
+
raise ValueError(f"cleanup_grace must be >= 0, got {cleanup_grace!r}")
|
|
497
|
+
self.cleanup_grace = cleanup_grace
|
|
498
|
+
if use_signals and not _HAS_SETITIMER:
|
|
499
|
+
raise RuntimeError("use_signals=True needs signal.setitimer, unavailable here")
|
|
500
|
+
self.seconds = seconds
|
|
501
|
+
self.use_signals = use_signals
|
|
502
|
+
self._alarms: list[_Alarm] = []
|
|
503
|
+
self._timers: dict[asyncio.Task[Any], list[Any]] = {}
|
|
504
|
+
|
|
505
|
+
def _limit(self, label: str) -> float:
|
|
506
|
+
"""The limit for a call starting now: ``seconds`` clamped to the deadline."""
|
|
507
|
+
left = _deadline.remaining()
|
|
508
|
+
if left is None or left >= self.seconds:
|
|
509
|
+
return self.seconds
|
|
510
|
+
if left <= 0:
|
|
511
|
+
_core.emit("timeout.expired", label, seconds=0.0)
|
|
512
|
+
raise TimeoutExpired(f"{label} not started: the deadline has passed", seconds=0.0)
|
|
513
|
+
return left
|
|
514
|
+
|
|
515
|
+
# -- decorator ----------------------------------------------------------------------
|
|
516
|
+
|
|
517
|
+
def __call__(self, func: F, /) -> F:
|
|
518
|
+
"""Apply the limit to every call of ``func`` (sync or ``async def``)."""
|
|
519
|
+
label = _core.qualname(func)
|
|
520
|
+
use_signals = self.use_signals
|
|
521
|
+
grace = self.cleanup_grace
|
|
522
|
+
|
|
523
|
+
if grace is not None and not _core.is_async(func):
|
|
524
|
+
raise TypeError(
|
|
525
|
+
f"timeout(cleanup_grace=...) only supports async def functions; {label} is "
|
|
526
|
+
"synchronous (signals and worker threads cannot give cleanup a bounded grace)"
|
|
527
|
+
)
|
|
528
|
+
|
|
529
|
+
if _core.is_async(func):
|
|
530
|
+
|
|
531
|
+
@functools.wraps(func)
|
|
532
|
+
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
533
|
+
seconds = self._limit(label)
|
|
534
|
+
if grace is not None:
|
|
535
|
+
return await _await_with_grace(func(*args, **kwargs), label, seconds, grace)
|
|
536
|
+
return await _await_with_deadline(func(*args, **kwargs), label, seconds)
|
|
537
|
+
|
|
538
|
+
return cast(F, async_wrapper)
|
|
539
|
+
|
|
540
|
+
@functools.wraps(func)
|
|
541
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
542
|
+
if use_signals and not _signals_usable():
|
|
543
|
+
raise RuntimeError(
|
|
544
|
+
"use_signals=True: signal-based timeouts only work in the main thread"
|
|
545
|
+
)
|
|
546
|
+
seconds = self._limit(label)
|
|
547
|
+
if use_signals or (use_signals is None and _signals_usable()):
|
|
548
|
+
with _Alarm(label, seconds):
|
|
549
|
+
return func(*args, **kwargs)
|
|
550
|
+
return _call_in_worker(func, label, seconds, *args, **kwargs)
|
|
551
|
+
|
|
552
|
+
return cast(F, wrapper)
|
|
553
|
+
|
|
554
|
+
# -- sync context manager -----------------------------------------------------------
|
|
555
|
+
|
|
556
|
+
def __enter__(self) -> Timeout:
|
|
557
|
+
if self.cleanup_grace is not None:
|
|
558
|
+
raise TypeError("timeout(cleanup_grace=...) is only supported with `async with`")
|
|
559
|
+
if self.use_signals is False or not _signals_usable():
|
|
560
|
+
raise RuntimeError(
|
|
561
|
+
"`with timeout(...)` needs SIGALRM (Unix) and the main thread, because a "
|
|
562
|
+
"block of code cannot be moved to a worker thread; decorate a function "
|
|
563
|
+
"with @timeout instead, or use `async with timeout(...)` in async code"
|
|
564
|
+
)
|
|
565
|
+
alarm = _Alarm(_BLOCK, self._limit(_BLOCK))
|
|
566
|
+
alarm.__enter__()
|
|
567
|
+
self._alarms.append(alarm)
|
|
568
|
+
return self
|
|
569
|
+
|
|
570
|
+
def __exit__(
|
|
571
|
+
self,
|
|
572
|
+
exc_type: type[BaseException] | None,
|
|
573
|
+
exc: BaseException | None,
|
|
574
|
+
tb: TracebackType | None,
|
|
575
|
+
) -> None:
|
|
576
|
+
self._alarms.pop().__exit__(exc_type, exc, tb)
|
|
577
|
+
|
|
578
|
+
# -- async context manager ----------------------------------------------------------
|
|
579
|
+
|
|
580
|
+
async def __aenter__(self) -> Timeout:
|
|
581
|
+
task = asyncio.current_task()
|
|
582
|
+
if task is None: # pragma: no cover - asyncio always runs coroutines in a task
|
|
583
|
+
raise RuntimeError("`async with timeout(...)` must run inside an asyncio task")
|
|
584
|
+
seconds = self._limit(_BLOCK)
|
|
585
|
+
timer: Any = (
|
|
586
|
+
_async_timer(seconds)
|
|
587
|
+
if self.cleanup_grace is None
|
|
588
|
+
else _GraceTimer(seconds, self.cleanup_grace, _BLOCK)
|
|
589
|
+
)
|
|
590
|
+
await timer.__aenter__()
|
|
591
|
+
self._timers.setdefault(task, []).append((timer, seconds))
|
|
592
|
+
return self
|
|
593
|
+
|
|
594
|
+
async def __aexit__(
|
|
595
|
+
self,
|
|
596
|
+
exc_type: type[BaseException] | None,
|
|
597
|
+
exc: BaseException | None,
|
|
598
|
+
tb: TracebackType | None,
|
|
599
|
+
) -> None:
|
|
600
|
+
task = asyncio.current_task()
|
|
601
|
+
assert task is not None
|
|
602
|
+
stack = self._timers[task]
|
|
603
|
+
timer, seconds = stack.pop()
|
|
604
|
+
if not stack:
|
|
605
|
+
del self._timers[task]
|
|
606
|
+
try:
|
|
607
|
+
await timer.__aexit__(exc_type, exc, tb)
|
|
608
|
+
except TimeoutError as error:
|
|
609
|
+
if not timer.expired(): # pragma: no cover - only our own expiry is converted
|
|
610
|
+
raise
|
|
611
|
+
raise _expired(_BLOCK, seconds) from error
|
|
612
|
+
|
|
613
|
+
def __repr__(self) -> str:
|
|
614
|
+
return (
|
|
615
|
+
f"Timeout(seconds={self.seconds!r}, use_signals={self.use_signals!r}, "
|
|
616
|
+
f"cleanup_grace={self.cleanup_grace!r})"
|
|
617
|
+
)
|
|
618
|
+
|
|
619
|
+
|
|
620
|
+
def timeout(
|
|
621
|
+
seconds: float, *, use_signals: bool | None = None, cleanup_grace: float | None = None
|
|
622
|
+
) -> Timeout:
|
|
623
|
+
"""Raise `TimeoutExpired` if a call or block runs longer than ``seconds``.
|
|
624
|
+
|
|
625
|
+
``TimeoutExpired`` subclasses the built-in ``TimeoutError``. The returned `Timeout`
|
|
626
|
+
works three ways:
|
|
627
|
+
|
|
628
|
+
* ``@timeout(5)`` decorates sync and ``async def`` functions.
|
|
629
|
+
* ``async with timeout(5):`` limits a block in async code (`asyncio.timeout` on
|
|
630
|
+
Python 3.11+, a cancellation timer on 3.10).
|
|
631
|
+
* ``with timeout(5):`` limits a block in sync code, with signals only: it needs
|
|
632
|
+
``signal.setitimer`` (Unix) and the main thread, otherwise entering it raises
|
|
633
|
+
``RuntimeError`` (a block cannot be moved to a worker thread; decorate a function
|
|
634
|
+
instead).
|
|
635
|
+
|
|
636
|
+
How each kind of code is stopped:
|
|
637
|
+
|
|
638
|
+
* ``async def`` code is cancelled at the deadline. Only the timeout's own expiry is
|
|
639
|
+
turned into ``TimeoutExpired``: a ``TimeoutError`` raised by the wrapped code
|
|
640
|
+
itself propagates unchanged.
|
|
641
|
+
* Sync code, with signals (the default on Unix when called from the main thread):
|
|
642
|
+
a ``SIGALRM`` timer from `signal.setitimer` interrupts the call. The previous
|
|
643
|
+
``SIGALRM`` handler and any already-armed timer are restored afterwards, so nested
|
|
644
|
+
timeouts work.
|
|
645
|
+
* Sync functions, without signals (Windows, non-main threads, or
|
|
646
|
+
``use_signals=False``): the call runs in a daemon worker thread and the caller stops
|
|
647
|
+
waiting at the deadline. **This fallback cannot stop the work** — the function keeps
|
|
648
|
+
running in the background until it returns on its own, and its result is discarded.
|
|
649
|
+
Workers are pooled daemon threads: an idle one is reused for the next call (at most
|
|
650
|
+
8 are kept idle, each exits after about 5 s without work), a worker whose call timed
|
|
651
|
+
out is abandoned and never reused, and the caller's ``contextvars`` are copied into
|
|
652
|
+
the worker on every call.
|
|
653
|
+
|
|
654
|
+
Inside a `python_wrap.deadline.deadline` block the limit is clamped to the time the
|
|
655
|
+
deadline has left; if the deadline has already passed, the call (or block) is not
|
|
656
|
+
started and ``TimeoutExpired`` is raised at once. Every expiry emits the
|
|
657
|
+
``timeout.expired`` event.
|
|
658
|
+
|
|
659
|
+
Warning:
|
|
660
|
+
With the signal strategy, ``TimeoutExpired`` is raised *inside* the wrapped
|
|
661
|
+
code, at whatever line it is executing when the alarm fires. A broad
|
|
662
|
+
``except Exception:`` (or ``except TimeoutError:``) in that code catches it like
|
|
663
|
+
any other error; if the code then carries on, the time limit is silently lost and
|
|
664
|
+
the call returns normally. ``TimeoutExpired`` deliberately stays an ``Exception``
|
|
665
|
+
(it subclasses ``TimeoutError``), so re-raise it explicitly in such handlers::
|
|
666
|
+
|
|
667
|
+
except TimeoutExpired:
|
|
668
|
+
raise
|
|
669
|
+
except Exception:
|
|
670
|
+
...
|
|
671
|
+
|
|
672
|
+
The same applies to context managers that suppress exceptions and to
|
|
673
|
+
``finally`` blocks that ``return``. ``async def`` code has the analogous
|
|
674
|
+
caveat for code that swallows ``asyncio.CancelledError``.
|
|
675
|
+
|
|
676
|
+
Args:
|
|
677
|
+
seconds: Time limit in seconds; must be positive.
|
|
678
|
+
use_signals: ``True`` forces the signal strategy, ``False`` forces the thread
|
|
679
|
+
fallback (and so rules out ``with timeout(...)``), ``None`` (default) picks
|
|
680
|
+
signals when they are usable.
|
|
681
|
+
cleanup_grace: Async only. Seconds the code may spend in its ``finally`` /
|
|
682
|
+
``except CancelledError`` cleanup after the limit is hit. As a decorator, the
|
|
683
|
+
call runs as a child task: at the limit it is cancelled and given
|
|
684
|
+
``cleanup_grace`` more seconds to finish; then ``TimeoutExpired`` is raised
|
|
685
|
+
(even if the cleanup finished), and a task still running is abandoned with a
|
|
686
|
+
warning logged to the package logger (it keeps running until it ends on its own).
|
|
687
|
+
As ``async with timeout(5, cleanup_grace=1)`` the block cannot be moved to
|
|
688
|
+
another task, so after the grace period the cleanup is cancelled a second time
|
|
689
|
+
(with a warning) rather than abandoned; cleanup that suppresses cancellation
|
|
690
|
+
cannot be stopped. Sync functions and ``with`` blocks are not supported and
|
|
691
|
+
raise ``TypeError`` (signals and worker threads cannot bound cleanup).
|
|
692
|
+
``None`` (default) keeps the plain behaviour. Real time: not affected by the
|
|
693
|
+
test overrides.
|
|
694
|
+
|
|
695
|
+
Returns:
|
|
696
|
+
A `Timeout`: a decorator that preserves the wrapped function's type, and a sync
|
|
697
|
+
and async context manager.
|
|
698
|
+
|
|
699
|
+
Raises:
|
|
700
|
+
ValueError: If ``seconds`` is not positive or ``cleanup_grace`` is negative.
|
|
701
|
+
TypeError: If ``cleanup_grace`` is set and the target is a sync function or a
|
|
702
|
+
plain ``with`` block.
|
|
703
|
+
RuntimeError: If ``use_signals=True`` but ``signal.setitimer`` does not exist on
|
|
704
|
+
this platform (at creation time); if ``use_signals=True`` and the call is not
|
|
705
|
+
made from the main thread (at call time); or on entering ``with`` where
|
|
706
|
+
signals cannot be used.
|
|
707
|
+
TimeoutExpired: From the wrapped call or block, when the time limit is exceeded.
|
|
708
|
+
|
|
709
|
+
Example:
|
|
710
|
+
```python
|
|
711
|
+
import asyncio
|
|
712
|
+
from python_wrap.concurrency import timeout
|
|
713
|
+
from python_wrap.errors import TimeoutExpired
|
|
714
|
+
|
|
715
|
+
@timeout(5)
|
|
716
|
+
def fetch() -> str:
|
|
717
|
+
return "done"
|
|
718
|
+
|
|
719
|
+
assert fetch() == "done"
|
|
720
|
+
|
|
721
|
+
async def main() -> None:
|
|
722
|
+
try:
|
|
723
|
+
async with timeout(0.01):
|
|
724
|
+
await asyncio.sleep(1)
|
|
725
|
+
except TimeoutExpired as exc:
|
|
726
|
+
assert exc.seconds == 0.01
|
|
727
|
+
|
|
728
|
+
asyncio.run(main())
|
|
729
|
+
```
|
|
730
|
+
"""
|
|
731
|
+
return Timeout(seconds, use_signals=use_signals, cleanup_grace=cleanup_grace)
|
|
732
|
+
|
|
733
|
+
|
|
734
|
+
# ---------------------------------------------------------------------------------------
|
|
735
|
+
# rate_limit
|
|
736
|
+
# ---------------------------------------------------------------------------------------
|
|
737
|
+
|
|
738
|
+
_SWEEP_EVERY = 1024
|
|
739
|
+
|
|
740
|
+
Limit = tuple[int, float]
|
|
741
|
+
"""One rate limit: ``(calls, period_seconds)``."""
|
|
742
|
+
|
|
743
|
+
|
|
744
|
+
def _limits(
|
|
745
|
+
calls: int | None, period: float | None, limits: Iterable[Limit] | None
|
|
746
|
+
) -> tuple[Limit, ...]:
|
|
747
|
+
"""Validate and combine the ``calls``/``period`` pair and ``limits``."""
|
|
748
|
+
result: list[Limit] = []
|
|
749
|
+
if calls is not None or period is not None:
|
|
750
|
+
if calls is None or period is None:
|
|
751
|
+
raise TypeError("rate limit: give both calls and period")
|
|
752
|
+
result.append((calls, period))
|
|
753
|
+
if limits is not None:
|
|
754
|
+
result.extend((c, p) for c, p in limits)
|
|
755
|
+
if not result:
|
|
756
|
+
raise TypeError("rate limit: give calls and period, or limits=[(calls, period), ...]")
|
|
757
|
+
for c, p in result:
|
|
758
|
+
if c < 1:
|
|
759
|
+
raise ValueError(f"rate_limit calls must be >= 1, got {c!r}")
|
|
760
|
+
if p <= 0:
|
|
761
|
+
raise ValueError(f"rate_limit period must be positive, got {p!r}")
|
|
762
|
+
return tuple(result)
|
|
763
|
+
|
|
764
|
+
|
|
765
|
+
def _describe(limits: tuple[Limit, ...]) -> str:
|
|
766
|
+
return " and ".join(f"{c} calls per {p:g}s" for c, p in limits)
|
|
767
|
+
|
|
768
|
+
|
|
769
|
+
class _SlidingWindow:
|
|
770
|
+
"""Thread-safe sliding-window logs of call timestamps: per key, one deque per limit."""
|
|
771
|
+
|
|
772
|
+
def __init__(self, limits: tuple[Limit, ...]) -> None:
|
|
773
|
+
self.limits = limits
|
|
774
|
+
self._lock = threading.Lock()
|
|
775
|
+
self._buckets: dict[Hashable, tuple[deque[float], ...]] = {}
|
|
776
|
+
self._since_sweep = 0
|
|
777
|
+
|
|
778
|
+
def acquire(self, key: Hashable) -> float:
|
|
779
|
+
"""Record a call and return 0.0, or return how long to wait before retrying.
|
|
780
|
+
|
|
781
|
+
A call is recorded only if *every* limit allows it.
|
|
782
|
+
"""
|
|
783
|
+
with self._lock:
|
|
784
|
+
# Read the clock under the lock: a call that queued for the lock must be
|
|
785
|
+
# recorded at the time it is admitted, not the time it arrived.
|
|
786
|
+
now = _core.monotonic()
|
|
787
|
+
self._since_sweep += 1
|
|
788
|
+
if self._since_sweep >= _SWEEP_EVERY:
|
|
789
|
+
self._sweep(now)
|
|
790
|
+
buckets = self._buckets.get(key)
|
|
791
|
+
if buckets is None:
|
|
792
|
+
buckets = self._buckets[key] = tuple(deque() for _ in self.limits)
|
|
793
|
+
wait = 0.0
|
|
794
|
+
for (calls, period), bucket in zip(self.limits, buckets, strict=True):
|
|
795
|
+
horizon = now - period
|
|
796
|
+
while bucket and bucket[0] <= horizon:
|
|
797
|
+
bucket.popleft()
|
|
798
|
+
if len(bucket) >= calls:
|
|
799
|
+
wait = max(wait, bucket[0] + period - now)
|
|
800
|
+
if wait > 0:
|
|
801
|
+
return wait
|
|
802
|
+
for bucket in buckets:
|
|
803
|
+
bucket.append(now)
|
|
804
|
+
return 0.0
|
|
805
|
+
|
|
806
|
+
def admit(self, key: Hashable) -> None:
|
|
807
|
+
"""Record a call now even if a limit is full (for `fast_mode()`'s blocked calls)."""
|
|
808
|
+
with self._lock:
|
|
809
|
+
now = _core.monotonic()
|
|
810
|
+
buckets = self._buckets.get(key)
|
|
811
|
+
if buckets is None:
|
|
812
|
+
buckets = self._buckets[key] = tuple(deque() for _ in self.limits)
|
|
813
|
+
for bucket in buckets:
|
|
814
|
+
bucket.append(now)
|
|
815
|
+
|
|
816
|
+
def clear(self) -> None:
|
|
817
|
+
with self._lock:
|
|
818
|
+
self._buckets.clear()
|
|
819
|
+
self._since_sweep = 0
|
|
820
|
+
|
|
821
|
+
def _sweep(self, now: float) -> None:
|
|
822
|
+
self._since_sweep = 0
|
|
823
|
+
periods = [period for _, period in self.limits]
|
|
824
|
+
|
|
825
|
+
def is_stale(buckets: tuple[deque[float], ...]) -> bool:
|
|
826
|
+
return all(not b or b[-1] <= now - p for p, b in zip(periods, buckets, strict=True))
|
|
827
|
+
|
|
828
|
+
stale = [k for k, buckets in self._buckets.items() if is_stale(buckets)]
|
|
829
|
+
for k in stale:
|
|
830
|
+
del self._buckets[k]
|
|
831
|
+
|
|
832
|
+
|
|
833
|
+
class RateLimiter:
|
|
834
|
+
"""A reusable sliding-window rate limit, shared by everything that uses it.
|
|
835
|
+
|
|
836
|
+
Use one instance to give several functions or code blocks a single budget:
|
|
837
|
+
|
|
838
|
+
* as a decorator (on sync or ``async def`` functions and methods);
|
|
839
|
+
* as ``with limiter:`` / ``async with limiter:`` around a block;
|
|
840
|
+
* through ``acquire()`` / ``try_acquire()`` / ``acquire_async()``.
|
|
841
|
+
|
|
842
|
+
A call is allowed when fewer than ``calls`` calls started in the last ``period``
|
|
843
|
+
seconds (a sliding-window log, so no ``2 * calls`` burst at a window boundary). With
|
|
844
|
+
``limits=[(5, 1.0), (600, 900.0)]`` every limit must allow the call. Rejected or
|
|
845
|
+
still-waiting calls are not counted. Thread-safe; async callers wait with ``await``
|
|
846
|
+
and never block the event loop. While ``python_wrap.testing`` disables rate limits,
|
|
847
|
+
every call goes through and nothing is recorded. Under ``fast_mode()`` (``no_sleep``)
|
|
848
|
+
a blocking call that would wait is let through at once and still recorded, so later
|
|
849
|
+
non-blocking calls see the limit; under ``fake_clock()`` the wait advances the fake
|
|
850
|
+
clock instead.
|
|
851
|
+
|
|
852
|
+
Events: ``rate_limit.rejected(retry_after)`` for each rejected call and
|
|
853
|
+
``rate_limit.waited(seconds)`` after a call had to wait, reported under ``name``
|
|
854
|
+
(or the decorated function's qualified name).
|
|
855
|
+
|
|
856
|
+
Args:
|
|
857
|
+
calls: Maximum number of calls per ``period`` (>= 1).
|
|
858
|
+
period: Window length in seconds (> 0).
|
|
859
|
+
limits: More ``(calls, period)`` pairs; may replace or add to ``calls``/``period``.
|
|
860
|
+
block: What the decorator and ``with`` do when over the limit: ``False``
|
|
861
|
+
(default) raises `RateLimitExceeded`, ``True`` waits for a slot.
|
|
862
|
+
name: Label for error messages and events.
|
|
863
|
+
|
|
864
|
+
Raises:
|
|
865
|
+
TypeError: If no limit is given, or only one of ``calls`` and ``period``.
|
|
866
|
+
ValueError: If a limit is out of range.
|
|
867
|
+
|
|
868
|
+
Example:
|
|
869
|
+
```python
|
|
870
|
+
from python_wrap.concurrency import RateLimiter
|
|
871
|
+
|
|
872
|
+
github = RateLimiter(limits=[(4, 1.0), (5000, 3600.0)], name="github")
|
|
873
|
+
|
|
874
|
+
@github
|
|
875
|
+
def get_repo(name: str) -> str:
|
|
876
|
+
return name
|
|
877
|
+
|
|
878
|
+
@github
|
|
879
|
+
def get_user(name: str) -> str:
|
|
880
|
+
return name
|
|
881
|
+
|
|
882
|
+
assert get_repo("a") == "a" and get_user("b") == "b" # one shared budget
|
|
883
|
+
with github: # a plain block counts too
|
|
884
|
+
pass
|
|
885
|
+
assert github.try_acquire()
|
|
886
|
+
assert not github.try_acquire() # 4 calls in the last second
|
|
887
|
+
```
|
|
888
|
+
"""
|
|
889
|
+
|
|
890
|
+
def __reduce__(self) -> NoReturn:
|
|
891
|
+
_core.refuse_pickle(self)
|
|
892
|
+
|
|
893
|
+
def __init__(
|
|
894
|
+
self,
|
|
895
|
+
calls: int | None = None,
|
|
896
|
+
period: float | None = None,
|
|
897
|
+
*,
|
|
898
|
+
limits: Iterable[Limit] | None = None,
|
|
899
|
+
block: bool = False,
|
|
900
|
+
name: str | None = None,
|
|
901
|
+
) -> None:
|
|
902
|
+
self.limits = _limits(calls, period, limits)
|
|
903
|
+
self.block = block
|
|
904
|
+
self.name = name
|
|
905
|
+
self._window = _SlidingWindow(self.limits)
|
|
906
|
+
|
|
907
|
+
def __repr__(self) -> str:
|
|
908
|
+
return f"RateLimiter(name={self.name!r}, limits={list(self.limits)!r})"
|
|
909
|
+
|
|
910
|
+
# -- core ---------------------------------------------------------------------------
|
|
911
|
+
|
|
912
|
+
def _rejected(self, label: str, retry_after: float) -> RateLimitExceeded:
|
|
913
|
+
_core.emit("rate_limit.rejected", label, retry_after=retry_after)
|
|
914
|
+
return RateLimitExceeded(
|
|
915
|
+
f"{label} rate limit of {_describe(self.limits)} exceeded; "
|
|
916
|
+
f"retry after {retry_after:.3g}s",
|
|
917
|
+
retry_after=retry_after,
|
|
918
|
+
)
|
|
919
|
+
|
|
920
|
+
def _gives_up(self, label: str, wait: float, give_up: float | None) -> bool:
|
|
921
|
+
"""True (after emitting the rejection) if ``wait`` overshoots ``give_up``."""
|
|
922
|
+
if give_up is not None and _core.monotonic() + wait > give_up:
|
|
923
|
+
_core.emit("rate_limit.rejected", label, retry_after=wait)
|
|
924
|
+
return True
|
|
925
|
+
return False
|
|
926
|
+
|
|
927
|
+
def _take(
|
|
928
|
+
self, key: Hashable, label: str, *, block: bool, timeout: float | None, strict: bool
|
|
929
|
+
) -> bool:
|
|
930
|
+
"""Take a slot in a thread. ``strict`` raises on rejection instead of returning False."""
|
|
931
|
+
if _core.overrides.disable_rate_limits:
|
|
932
|
+
return True
|
|
933
|
+
wait = self._window.acquire(key)
|
|
934
|
+
if wait <= 0:
|
|
935
|
+
return True
|
|
936
|
+
if not block:
|
|
937
|
+
error = self._rejected(label, wait)
|
|
938
|
+
if strict:
|
|
939
|
+
raise error
|
|
940
|
+
return False
|
|
941
|
+
start = _core.monotonic()
|
|
942
|
+
give_up = None if timeout is None else start + timeout
|
|
943
|
+
while wait > 0:
|
|
944
|
+
if self._gives_up(label, wait, give_up):
|
|
945
|
+
return False
|
|
946
|
+
_core.sleep(wait)
|
|
947
|
+
if _core.overrides.disable_rate_limits:
|
|
948
|
+
break
|
|
949
|
+
wait = self._window.acquire(key)
|
|
950
|
+
if wait > 0 and _core.overrides.no_sleep:
|
|
951
|
+
# fast_mode() and the sleep did not move the clock: admit the call
|
|
952
|
+
# instead of spinning (it still counts toward the limit).
|
|
953
|
+
self._window.admit(key)
|
|
954
|
+
break
|
|
955
|
+
_core.emit("rate_limit.waited", label, seconds=_core.monotonic() - start)
|
|
956
|
+
return True
|
|
957
|
+
|
|
958
|
+
async def _atake(
|
|
959
|
+
self, key: Hashable, label: str, *, block: bool, timeout: float | None, strict: bool
|
|
960
|
+
) -> bool:
|
|
961
|
+
"""`_take` for coroutines: waits with ``await``."""
|
|
962
|
+
if _core.overrides.disable_rate_limits:
|
|
963
|
+
return True
|
|
964
|
+
wait = self._window.acquire(key)
|
|
965
|
+
if wait <= 0:
|
|
966
|
+
return True
|
|
967
|
+
if not block:
|
|
968
|
+
error = self._rejected(label, wait)
|
|
969
|
+
if strict:
|
|
970
|
+
raise error
|
|
971
|
+
return False
|
|
972
|
+
start = _core.monotonic()
|
|
973
|
+
give_up = None if timeout is None else start + timeout
|
|
974
|
+
while wait > 0:
|
|
975
|
+
if self._gives_up(label, wait, give_up):
|
|
976
|
+
return False
|
|
977
|
+
await _core.async_sleep(wait)
|
|
978
|
+
if _core.overrides.disable_rate_limits:
|
|
979
|
+
break
|
|
980
|
+
wait = self._window.acquire(key)
|
|
981
|
+
if wait > 0 and _core.overrides.no_sleep:
|
|
982
|
+
self._window.admit(key) # as in `_take`: never spin under fast_mode()
|
|
983
|
+
break
|
|
984
|
+
_core.emit("rate_limit.waited", label, seconds=_core.monotonic() - start)
|
|
985
|
+
return True
|
|
986
|
+
|
|
987
|
+
def _label(self) -> str:
|
|
988
|
+
return self.name or "rate_limiter"
|
|
989
|
+
|
|
990
|
+
# -- explicit API -------------------------------------------------------------------
|
|
991
|
+
|
|
992
|
+
def reset(self) -> None:
|
|
993
|
+
"""Forget every recorded call (for every key), as if the limiter were new.
|
|
994
|
+
|
|
995
|
+
Meant for tests: share one `RateLimiter` (or reach a `rate_limit` limiter with
|
|
996
|
+
`limiter_of`) and reset it between tests. Thread-safe; calls waiting for a slot
|
|
997
|
+
re-check and may then go through.
|
|
998
|
+
"""
|
|
999
|
+
self._window.clear()
|
|
1000
|
+
|
|
1001
|
+
def try_acquire(self) -> bool:
|
|
1002
|
+
"""Take a slot if one is free right now.
|
|
1003
|
+
|
|
1004
|
+
Returns:
|
|
1005
|
+
``True`` if the call may proceed (it is counted), ``False`` otherwise.
|
|
1006
|
+
"""
|
|
1007
|
+
return self._take(None, self._label(), block=False, timeout=None, strict=False)
|
|
1008
|
+
|
|
1009
|
+
def acquire(self, block: bool = True, timeout: float | None = None) -> bool:
|
|
1010
|
+
"""Take a slot, waiting for one in this thread if ``block`` is true.
|
|
1011
|
+
|
|
1012
|
+
Never call it from a coroutine; use `acquire_async` there.
|
|
1013
|
+
|
|
1014
|
+
Args:
|
|
1015
|
+
block: Wait for a slot (``True``) or return at once (``False``).
|
|
1016
|
+
timeout: Give up after this many seconds (``None`` waits as long as needed).
|
|
1017
|
+
Slots only free up as time passes, so the call returns ``False`` right
|
|
1018
|
+
away when the next slot is further away than ``timeout``.
|
|
1019
|
+
|
|
1020
|
+
Returns:
|
|
1021
|
+
``True`` if a slot was taken, ``False`` if not.
|
|
1022
|
+
"""
|
|
1023
|
+
return self._take(None, self._label(), block=block, timeout=timeout, strict=False)
|
|
1024
|
+
|
|
1025
|
+
async def acquire_async(self, block: bool = True, timeout: float | None = None) -> bool:
|
|
1026
|
+
"""`acquire` for async code: waits with ``await``, never blocking the loop.
|
|
1027
|
+
|
|
1028
|
+
Args:
|
|
1029
|
+
block: Wait for a slot (``True``) or return at once (``False``).
|
|
1030
|
+
timeout: Give up after this many seconds (``None`` waits as long as needed).
|
|
1031
|
+
|
|
1032
|
+
Returns:
|
|
1033
|
+
``True`` if a slot was taken, ``False`` if not.
|
|
1034
|
+
"""
|
|
1035
|
+
return await self._atake(None, self._label(), block=block, timeout=timeout, strict=False)
|
|
1036
|
+
|
|
1037
|
+
def __enter__(self) -> RateLimiter:
|
|
1038
|
+
self._take(None, self._label(), block=self.block, timeout=None, strict=True)
|
|
1039
|
+
return self
|
|
1040
|
+
|
|
1041
|
+
def __exit__(
|
|
1042
|
+
self,
|
|
1043
|
+
exc_type: type[BaseException] | None,
|
|
1044
|
+
exc: BaseException | None,
|
|
1045
|
+
tb: TracebackType | None,
|
|
1046
|
+
) -> None:
|
|
1047
|
+
return None
|
|
1048
|
+
|
|
1049
|
+
async def __aenter__(self) -> RateLimiter:
|
|
1050
|
+
await self._atake(None, self._label(), block=self.block, timeout=None, strict=True)
|
|
1051
|
+
return self
|
|
1052
|
+
|
|
1053
|
+
async def __aexit__(
|
|
1054
|
+
self,
|
|
1055
|
+
exc_type: type[BaseException] | None,
|
|
1056
|
+
exc: BaseException | None,
|
|
1057
|
+
tb: TracebackType | None,
|
|
1058
|
+
) -> None:
|
|
1059
|
+
return None
|
|
1060
|
+
|
|
1061
|
+
# -- decorator ----------------------------------------------------------------------
|
|
1062
|
+
|
|
1063
|
+
def __call__(self, func: F, /) -> F:
|
|
1064
|
+
"""Count every call of ``func`` against this limiter (sync or ``async def``)."""
|
|
1065
|
+
return self._wrap(func, None)
|
|
1066
|
+
|
|
1067
|
+
def _wrap(self, func: F, key: Callable[..., Hashable] | None) -> F:
|
|
1068
|
+
label = self.name or _core.qualname(func)
|
|
1069
|
+
block = self.block
|
|
1070
|
+
|
|
1071
|
+
if _core.is_async(func):
|
|
1072
|
+
|
|
1073
|
+
@functools.wraps(func)
|
|
1074
|
+
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
1075
|
+
k = None if key is None else key(*args, **kwargs)
|
|
1076
|
+
await self._atake(k, label, block=block, timeout=None, strict=True)
|
|
1077
|
+
return await func(*args, **kwargs)
|
|
1078
|
+
|
|
1079
|
+
async_wrapper.__rate_limiter__ = self # type: ignore[attr-defined]
|
|
1080
|
+
return cast(F, async_wrapper)
|
|
1081
|
+
|
|
1082
|
+
@functools.wraps(func)
|
|
1083
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
1084
|
+
k = None if key is None else key(*args, **kwargs)
|
|
1085
|
+
self._take(k, label, block=block, timeout=None, strict=True)
|
|
1086
|
+
return func(*args, **kwargs)
|
|
1087
|
+
|
|
1088
|
+
wrapper.__rate_limiter__ = self # type: ignore[attr-defined]
|
|
1089
|
+
return cast(F, wrapper)
|
|
1090
|
+
|
|
1091
|
+
|
|
1092
|
+
def limiter_of(func: Callable[..., Any], /) -> RateLimiter | None:
|
|
1093
|
+
"""The `RateLimiter` behind a function decorated with `rate_limit` (or a limiter).
|
|
1094
|
+
|
|
1095
|
+
`rate_limit` keeps the decorated function's exact type, so its limiter is reached
|
|
1096
|
+
through this function rather than an attribute. Decorators applied on top with
|
|
1097
|
+
`functools.wraps` keep it reachable; bound methods work too.
|
|
1098
|
+
|
|
1099
|
+
Args:
|
|
1100
|
+
func: A function decorated with `rate_limit` or with a `RateLimiter`.
|
|
1101
|
+
|
|
1102
|
+
Returns:
|
|
1103
|
+
Its limiter (the outermost one, if several), or `None` if it has none.
|
|
1104
|
+
|
|
1105
|
+
Example:
|
|
1106
|
+
```python
|
|
1107
|
+
from python_wrap.concurrency import limiter_of, rate_limit
|
|
1108
|
+
|
|
1109
|
+
@rate_limit(1, 60.0)
|
|
1110
|
+
def ping() -> str:
|
|
1111
|
+
return "pong"
|
|
1112
|
+
|
|
1113
|
+
assert ping() == "pong"
|
|
1114
|
+
limiter = limiter_of(ping)
|
|
1115
|
+
assert limiter is not None and not limiter.try_acquire()
|
|
1116
|
+
limiter.reset() # e.g. between tests
|
|
1117
|
+
assert ping() == "pong"
|
|
1118
|
+
```
|
|
1119
|
+
"""
|
|
1120
|
+
limiter = getattr(func, "__rate_limiter__", None)
|
|
1121
|
+
return limiter if isinstance(limiter, RateLimiter) else None
|
|
1122
|
+
|
|
1123
|
+
|
|
1124
|
+
def rate_limit(
|
|
1125
|
+
calls: int | None = None,
|
|
1126
|
+
period: float | None = None,
|
|
1127
|
+
*,
|
|
1128
|
+
block: bool = False,
|
|
1129
|
+
key: Callable[..., Hashable] | None = None,
|
|
1130
|
+
limits: Iterable[Limit] | None = None,
|
|
1131
|
+
) -> Callable[[F], F]:
|
|
1132
|
+
"""Allow at most ``calls`` calls per ``period`` seconds.
|
|
1133
|
+
|
|
1134
|
+
Uses a sliding-window log: a call is allowed when fewer than ``calls`` calls started
|
|
1135
|
+
in the last ``period`` seconds. Unlike a fixed window, this never lets a burst of
|
|
1136
|
+
``2 * calls`` through at a window boundary. Rejected (or waiting) calls are not
|
|
1137
|
+
counted. The limiter is thread-safe and is shared by every caller of the decorated
|
|
1138
|
+
function — including every instance, for methods — unless ``key`` splits it.
|
|
1139
|
+
|
|
1140
|
+
Each decorated function gets its own `RateLimiter`, reachable with
|
|
1141
|
+
``limiter_of(func)`` (e.g. to ``.reset()`` it between tests). To share one budget
|
|
1142
|
+
between several functions, create a `RateLimiter` and use it as the decorator. Pass
|
|
1143
|
+
``limits=[(5, 1.0), (600, 900.0)]`` to enforce several limits at once (all must
|
|
1144
|
+
allow a call). While ``python_wrap.testing`` disables rate limits, every call goes
|
|
1145
|
+
through; under ``fast_mode()`` a ``block=True`` call that would wait goes through at
|
|
1146
|
+
once (and is counted). Emits ``rate_limit.rejected`` and ``rate_limit.waited``.
|
|
1147
|
+
|
|
1148
|
+
Args:
|
|
1149
|
+
calls: Maximum number of calls per window; at least 1.
|
|
1150
|
+
period: Window length in seconds; must be positive.
|
|
1151
|
+
block: If ``False`` (default) an over-limit call raises ``RateLimitExceeded``.
|
|
1152
|
+
If ``True`` it waits until allowed (``async def`` functions wait with
|
|
1153
|
+
``await``, never blocking the event loop).
|
|
1154
|
+
key: Optional function called with the call's arguments (including ``self`` for
|
|
1155
|
+
methods) that returns a hashable bucket key, e.g. a user id. Each key gets
|
|
1156
|
+
its own limit.
|
|
1157
|
+
limits: Further ``(calls, period)`` pairs, alone or on top of ``calls``/``period``.
|
|
1158
|
+
|
|
1159
|
+
Returns:
|
|
1160
|
+
A decorator that preserves the wrapped function's signature.
|
|
1161
|
+
|
|
1162
|
+
Raises:
|
|
1163
|
+
TypeError: If no limit is given, or only one of ``calls`` and ``period``.
|
|
1164
|
+
ValueError: If ``calls < 1`` or ``period <= 0`` (for any limit).
|
|
1165
|
+
RateLimitExceeded: From the wrapped call when over the limit and ``block`` is
|
|
1166
|
+
false. Its ``retry_after`` attribute says how many seconds until a slot frees.
|
|
1167
|
+
|
|
1168
|
+
Example:
|
|
1169
|
+
>>> @rate_limit(2, 1.0, key=lambda user, *_: user)
|
|
1170
|
+
... def send(user: str, text: str) -> None: ...
|
|
1171
|
+
>>> @rate_limit(limits=[(5, 1.0), (600, 900.0)])
|
|
1172
|
+
... def strava_call() -> None: ...
|
|
1173
|
+
"""
|
|
1174
|
+
all_limits = _limits(calls, period, limits)
|
|
1175
|
+
|
|
1176
|
+
def decorate(func: F) -> F:
|
|
1177
|
+
return RateLimiter(limits=all_limits, block=block)._wrap(func, key)
|
|
1178
|
+
|
|
1179
|
+
return decorate
|
|
1180
|
+
|
|
1181
|
+
|
|
1182
|
+
# ---------------------------------------------------------------------------------------
|
|
1183
|
+
# run_in_thread
|
|
1184
|
+
# ---------------------------------------------------------------------------------------
|
|
1185
|
+
|
|
1186
|
+
|
|
1187
|
+
class _RunInThreadDecorator(Protocol):
|
|
1188
|
+
@overload
|
|
1189
|
+
def __call__(
|
|
1190
|
+
self, func: Callable[P, Coroutine[Any, Any, T]], /
|
|
1191
|
+
) -> Callable[P, concurrent.futures.Future[T]]: ...
|
|
1192
|
+
@overload
|
|
1193
|
+
def __call__(self, func: Callable[P, T], /) -> Callable[P, concurrent.futures.Future[T]]: ...
|
|
1194
|
+
|
|
1195
|
+
|
|
1196
|
+
@overload
|
|
1197
|
+
def run_in_thread(
|
|
1198
|
+
func: Callable[P, Coroutine[Any, Any, T]], /
|
|
1199
|
+
) -> Callable[P, concurrent.futures.Future[T]]: ...
|
|
1200
|
+
@overload
|
|
1201
|
+
def run_in_thread(func: Callable[P, T], /) -> Callable[P, concurrent.futures.Future[T]]: ...
|
|
1202
|
+
@overload
|
|
1203
|
+
def run_in_thread(
|
|
1204
|
+
func: None = None, /, *, daemon: bool = False, name: str | None = None
|
|
1205
|
+
) -> _RunInThreadDecorator: ...
|
|
1206
|
+
def run_in_thread(
|
|
1207
|
+
func: Callable[..., Any] | None = None,
|
|
1208
|
+
/,
|
|
1209
|
+
*,
|
|
1210
|
+
daemon: bool = False,
|
|
1211
|
+
name: str | None = None,
|
|
1212
|
+
) -> Any:
|
|
1213
|
+
"""Run each call in a new thread and return a `concurrent.futures.Future`.
|
|
1214
|
+
|
|
1215
|
+
The future resolves to the function's return value, or raises its exception from
|
|
1216
|
+
``future.result()``. For ``async def`` functions the coroutine is run with
|
|
1217
|
+
`asyncio.run` inside the new thread (so it gets its own event loop), and the
|
|
1218
|
+
wrapper itself is a plain function returning a ``Future``. The caller's
|
|
1219
|
+
``contextvars`` are copied into the thread.
|
|
1220
|
+
|
|
1221
|
+
Changed in 0.2.0: the wrapper used to return the ``threading.Thread``, with no way to
|
|
1222
|
+
get the result. Call ``.result()`` on the returned future to wait (the equivalent of
|
|
1223
|
+
the old ``thread.join()``).
|
|
1224
|
+
|
|
1225
|
+
Args:
|
|
1226
|
+
func: The function to wrap (bare ``@run_in_thread`` form).
|
|
1227
|
+
daemon: Start daemon threads, which do not keep the interpreter alive.
|
|
1228
|
+
name: Thread name; defaults to ``run_in_thread-<qualname>``.
|
|
1229
|
+
|
|
1230
|
+
Returns:
|
|
1231
|
+
The wrapped function, or a decorator when called with options only.
|
|
1232
|
+
|
|
1233
|
+
Example:
|
|
1234
|
+
>>> @run_in_thread
|
|
1235
|
+
... def add(a: int, b: int) -> int:
|
|
1236
|
+
... return a + b
|
|
1237
|
+
>>> add(1, 2).result()
|
|
1238
|
+
3
|
|
1239
|
+
"""
|
|
1240
|
+
|
|
1241
|
+
def decorate(f: Callable[..., Any]) -> Callable[..., concurrent.futures.Future[Any]]:
|
|
1242
|
+
thread_name = name or f"run_in_thread-{_core.qualname(f)}"
|
|
1243
|
+
is_coro = _core.is_async(f)
|
|
1244
|
+
|
|
1245
|
+
@functools.wraps(f)
|
|
1246
|
+
def wrapper(*args: Any, **kwargs: Any) -> concurrent.futures.Future[Any]:
|
|
1247
|
+
future: concurrent.futures.Future[Any] = concurrent.futures.Future()
|
|
1248
|
+
context = contextvars.copy_context()
|
|
1249
|
+
|
|
1250
|
+
def target() -> Any:
|
|
1251
|
+
if is_coro:
|
|
1252
|
+
return asyncio.run(f(*args, **kwargs))
|
|
1253
|
+
return f(*args, **kwargs)
|
|
1254
|
+
|
|
1255
|
+
def runner() -> None:
|
|
1256
|
+
if not future.set_running_or_notify_cancel():
|
|
1257
|
+
return
|
|
1258
|
+
try:
|
|
1259
|
+
future.set_result(context.run(target))
|
|
1260
|
+
except BaseException as exc:
|
|
1261
|
+
future.set_exception(exc)
|
|
1262
|
+
|
|
1263
|
+
threading.Thread(target=runner, name=thread_name, daemon=daemon).start()
|
|
1264
|
+
return future
|
|
1265
|
+
|
|
1266
|
+
return wrapper
|
|
1267
|
+
|
|
1268
|
+
return decorate if func is None else decorate(func)
|
|
1269
|
+
|
|
1270
|
+
|
|
1271
|
+
# ---------------------------------------------------------------------------------------
|
|
1272
|
+
# throttle
|
|
1273
|
+
# ---------------------------------------------------------------------------------------
|
|
1274
|
+
|
|
1275
|
+
|
|
1276
|
+
class _ThrottleDecorator(Protocol):
|
|
1277
|
+
@overload
|
|
1278
|
+
def __call__(
|
|
1279
|
+
self, func: Callable[P, Coroutine[Any, Any, T]], /
|
|
1280
|
+
) -> Callable[P, Coroutine[Any, Any, T | None]]: ...
|
|
1281
|
+
@overload
|
|
1282
|
+
def __call__(self, func: Callable[P, T], /) -> Callable[P, T | None]: ...
|
|
1283
|
+
|
|
1284
|
+
|
|
1285
|
+
def throttle(interval: float) -> _ThrottleDecorator:
|
|
1286
|
+
"""Run the function at most once per ``interval`` seconds.
|
|
1287
|
+
|
|
1288
|
+
The first call runs and opens a window of ``interval`` seconds (measured from when
|
|
1289
|
+
the call started). Calls inside the window do not run; they return the result of the
|
|
1290
|
+
most recent run that *returned* (``None`` if none has returned yet). A call that
|
|
1291
|
+
raises propagates its exception to its own caller and still uses up its window, but
|
|
1292
|
+
does not replace the remembered result: calls skipped inside that window get the
|
|
1293
|
+
previous successful result (or ``None``), not the error. Thread-safe; the function
|
|
1294
|
+
itself runs outside the lock, so a slow call does not block skipped callers.
|
|
1295
|
+
``async def`` functions get the same semantics with an ``async`` wrapper.
|
|
1296
|
+
|
|
1297
|
+
Args:
|
|
1298
|
+
interval: Minimum seconds between runs; must be positive.
|
|
1299
|
+
|
|
1300
|
+
Returns:
|
|
1301
|
+
A decorator; the wrapped function returns ``R | None``.
|
|
1302
|
+
|
|
1303
|
+
Raises:
|
|
1304
|
+
ValueError: If ``interval`` is not positive.
|
|
1305
|
+
|
|
1306
|
+
Example:
|
|
1307
|
+
>>> @throttle(1.0)
|
|
1308
|
+
... def refresh() -> str:
|
|
1309
|
+
... return "fresh"
|
|
1310
|
+
"""
|
|
1311
|
+
if interval <= 0:
|
|
1312
|
+
raise ValueError(f"throttle interval must be positive, got {interval!r}")
|
|
1313
|
+
|
|
1314
|
+
def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
1315
|
+
lock = threading.Lock()
|
|
1316
|
+
last_start: float | None = None
|
|
1317
|
+
last_result: Any = None
|
|
1318
|
+
|
|
1319
|
+
name = _core.qualname(func)
|
|
1320
|
+
|
|
1321
|
+
def should_run() -> tuple[bool, Any]:
|
|
1322
|
+
nonlocal last_start
|
|
1323
|
+
if _core.overrides.disable_rate_limits:
|
|
1324
|
+
return True, None
|
|
1325
|
+
now = _core.monotonic()
|
|
1326
|
+
with lock:
|
|
1327
|
+
if last_start is None or now - last_start >= interval:
|
|
1328
|
+
last_start = now
|
|
1329
|
+
return True, None
|
|
1330
|
+
previous = last_result
|
|
1331
|
+
_core.emit("throttle.skipped", name)
|
|
1332
|
+
return False, previous
|
|
1333
|
+
|
|
1334
|
+
def store(result: Any) -> Any:
|
|
1335
|
+
nonlocal last_result
|
|
1336
|
+
with lock:
|
|
1337
|
+
last_result = result
|
|
1338
|
+
return result
|
|
1339
|
+
|
|
1340
|
+
if _core.is_async(func):
|
|
1341
|
+
|
|
1342
|
+
@functools.wraps(func)
|
|
1343
|
+
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
1344
|
+
run, previous = should_run()
|
|
1345
|
+
if not run:
|
|
1346
|
+
return previous
|
|
1347
|
+
return store(await func(*args, **kwargs))
|
|
1348
|
+
|
|
1349
|
+
return async_wrapper
|
|
1350
|
+
|
|
1351
|
+
@functools.wraps(func)
|
|
1352
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
1353
|
+
run, previous = should_run()
|
|
1354
|
+
if not run:
|
|
1355
|
+
return previous
|
|
1356
|
+
return store(func(*args, **kwargs))
|
|
1357
|
+
|
|
1358
|
+
return wrapper
|
|
1359
|
+
|
|
1360
|
+
return cast(_ThrottleDecorator, decorate)
|
|
1361
|
+
|
|
1362
|
+
|
|
1363
|
+
# ---------------------------------------------------------------------------------------
|
|
1364
|
+
# debounce
|
|
1365
|
+
# ---------------------------------------------------------------------------------------
|
|
1366
|
+
|
|
1367
|
+
|
|
1368
|
+
class Debounced(Protocol[P, R_co]):
|
|
1369
|
+
"""A debounced sync function: calls return ``None``; the real call happens later.
|
|
1370
|
+
|
|
1371
|
+
Binds ``self`` and ``cls`` like `python_wrap.caching.Memoized`, with the same typing
|
|
1372
|
+
limitations.
|
|
1373
|
+
"""
|
|
1374
|
+
|
|
1375
|
+
def __call__(self, *args: P.args, **kwargs: P.kwargs) -> None: ...
|
|
1376
|
+
|
|
1377
|
+
@overload
|
|
1378
|
+
def __get__(
|
|
1379
|
+
self: Debounced[Concatenate[type[T], P2], R_co], instance: None, owner: type[T]
|
|
1380
|
+
) -> _DebouncedOnClass[P, P2, R_co]: ...
|
|
1381
|
+
@overload
|
|
1382
|
+
def __get__(self, instance: None, owner: type | None = None) -> Debounced[P, R_co]: ...
|
|
1383
|
+
@overload
|
|
1384
|
+
def __get__(
|
|
1385
|
+
self: Debounced[Concatenate[Any, P2], R_co], instance: object, owner: type | None = None
|
|
1386
|
+
) -> Debounced[P2, R_co]: ...
|
|
1387
|
+
|
|
1388
|
+
def cancel(self) -> None:
|
|
1389
|
+
"""Drop the pending call, if any."""
|
|
1390
|
+
|
|
1391
|
+
def flush(self) -> R_co | None:
|
|
1392
|
+
"""Run the pending call now and return its result (``None`` if nothing pending)."""
|
|
1393
|
+
|
|
1394
|
+
|
|
1395
|
+
class _DebouncedOnClass(Debounced[P, R_co], Protocol[P, P2, R_co]):
|
|
1396
|
+
"""`Debounced` read from a class: accepts the call with and without the first argument.
|
|
1397
|
+
|
|
1398
|
+
Returned for class access (``Cls.attr``) when the first parameter could be ``cls``.
|
|
1399
|
+
pyright only picks it under ``@classmethod``; mypy cannot tell a classmethod from an
|
|
1400
|
+
instance or static method there, so both call forms are accepted.
|
|
1401
|
+
"""
|
|
1402
|
+
|
|
1403
|
+
@overload
|
|
1404
|
+
def __call__(self, *args: P2.args, **kwargs: P2.kwargs) -> None: ...
|
|
1405
|
+
@overload
|
|
1406
|
+
def __call__(self, *args: P.args, **kwargs: P.kwargs) -> None: ...
|
|
1407
|
+
|
|
1408
|
+
|
|
1409
|
+
class AsyncDebounced(Protocol[P, R_co]):
|
|
1410
|
+
"""A debounced ``async def`` function; see `debounce` for the await semantics.
|
|
1411
|
+
|
|
1412
|
+
Binds ``self`` and ``cls`` like `python_wrap.caching.Memoized`, with the same typing
|
|
1413
|
+
limitations.
|
|
1414
|
+
"""
|
|
1415
|
+
|
|
1416
|
+
def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R_co]: ...
|
|
1417
|
+
|
|
1418
|
+
@overload
|
|
1419
|
+
def __get__(
|
|
1420
|
+
self: AsyncDebounced[Concatenate[type[T], P2], R_co], instance: None, owner: type[T]
|
|
1421
|
+
) -> _AsyncDebouncedOnClass[P, P2, R_co]: ...
|
|
1422
|
+
@overload
|
|
1423
|
+
def __get__(self, instance: None, owner: type | None = None) -> AsyncDebounced[P, R_co]: ...
|
|
1424
|
+
@overload
|
|
1425
|
+
def __get__(
|
|
1426
|
+
self: AsyncDebounced[Concatenate[Any, P2], R_co],
|
|
1427
|
+
instance: object,
|
|
1428
|
+
owner: type | None = None,
|
|
1429
|
+
) -> AsyncDebounced[P2, R_co]: ...
|
|
1430
|
+
|
|
1431
|
+
def cancel(self) -> None:
|
|
1432
|
+
"""Drop the pending call; callers awaiting it get ``CancelledError``."""
|
|
1433
|
+
|
|
1434
|
+
def flush(self) -> Coroutine[Any, Any, R_co | None]:
|
|
1435
|
+
"""Run the pending call now; awaiting returns its result (``None`` if none)."""
|
|
1436
|
+
|
|
1437
|
+
|
|
1438
|
+
class _AsyncDebouncedOnClass(AsyncDebounced[P, R_co], Protocol[P, P2, R_co]):
|
|
1439
|
+
"""`AsyncDebounced` read from a class: accepts the call with and without the first argument.
|
|
1440
|
+
|
|
1441
|
+
Returned for class access (``Cls.attr``) when the first parameter could be ``cls``.
|
|
1442
|
+
pyright only picks it under ``@classmethod``; mypy cannot tell a classmethod from an
|
|
1443
|
+
instance or static method there, so both call forms are accepted.
|
|
1444
|
+
"""
|
|
1445
|
+
|
|
1446
|
+
@overload
|
|
1447
|
+
def __call__(self, *args: P2.args, **kwargs: P2.kwargs) -> Coroutine[Any, Any, R_co]: ...
|
|
1448
|
+
@overload
|
|
1449
|
+
def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R_co]: ...
|
|
1450
|
+
|
|
1451
|
+
|
|
1452
|
+
class _DebounceDecorator(Protocol):
|
|
1453
|
+
@overload
|
|
1454
|
+
def __call__( # type: ignore[overload-overlap]
|
|
1455
|
+
self, func: Callable[P, Coroutine[Any, Any, T]], /
|
|
1456
|
+
) -> AsyncDebounced[P, T]: ...
|
|
1457
|
+
@overload
|
|
1458
|
+
def __call__(self, func: Callable[P, T], /) -> Debounced[P, T]: ...
|
|
1459
|
+
|
|
1460
|
+
|
|
1461
|
+
_Pending = tuple[tuple[Any, ...], dict[str, Any]]
|
|
1462
|
+
|
|
1463
|
+
|
|
1464
|
+
def _debounce_sync(func: Callable[..., Any], wait: float) -> Callable[..., None]:
|
|
1465
|
+
lock = threading.Lock()
|
|
1466
|
+
timer: threading.Timer | None = None
|
|
1467
|
+
pending: _Pending | None = None
|
|
1468
|
+
|
|
1469
|
+
def take() -> _Pending | None:
|
|
1470
|
+
nonlocal timer, pending
|
|
1471
|
+
if timer is not None:
|
|
1472
|
+
timer.cancel()
|
|
1473
|
+
call, timer, pending = pending, None, None
|
|
1474
|
+
return call
|
|
1475
|
+
|
|
1476
|
+
def fire(me: list[threading.Timer]) -> None:
|
|
1477
|
+
with lock:
|
|
1478
|
+
if timer is not me[0]: # superseded or cancelled meanwhile
|
|
1479
|
+
return
|
|
1480
|
+
call = take()
|
|
1481
|
+
if call is not None:
|
|
1482
|
+
func(*call[0], **call[1])
|
|
1483
|
+
|
|
1484
|
+
@functools.wraps(func)
|
|
1485
|
+
def wrapper(*args: Any, **kwargs: Any) -> None:
|
|
1486
|
+
nonlocal timer, pending
|
|
1487
|
+
with lock:
|
|
1488
|
+
take()
|
|
1489
|
+
me: list[threading.Timer] = []
|
|
1490
|
+
timer = threading.Timer(wait, fire, args=(me,))
|
|
1491
|
+
me.append(timer)
|
|
1492
|
+
timer.daemon = True
|
|
1493
|
+
pending = (args, kwargs)
|
|
1494
|
+
timer.start()
|
|
1495
|
+
|
|
1496
|
+
def cancel() -> None:
|
|
1497
|
+
with lock:
|
|
1498
|
+
take()
|
|
1499
|
+
|
|
1500
|
+
def flush() -> Any:
|
|
1501
|
+
with lock:
|
|
1502
|
+
call = take()
|
|
1503
|
+
return None if call is None else func(*call[0], **call[1])
|
|
1504
|
+
|
|
1505
|
+
wrapper.cancel = cancel # type: ignore[attr-defined]
|
|
1506
|
+
wrapper.flush = flush # type: ignore[attr-defined]
|
|
1507
|
+
return wrapper
|
|
1508
|
+
|
|
1509
|
+
|
|
1510
|
+
class _AsyncBurst:
|
|
1511
|
+
"""One burst of async debounced calls: the latest arguments and a shared outcome."""
|
|
1512
|
+
|
|
1513
|
+
def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
|
|
1514
|
+
self.outcome: asyncio.Future[Any] = loop.create_future()
|
|
1515
|
+
self.timer: asyncio.Task[None] | None = None
|
|
1516
|
+
self.call: _Pending | None = None
|
|
1517
|
+
|
|
1518
|
+
|
|
1519
|
+
def _debounce_async(func: Callable[..., Awaitable[Any]], wait: float) -> Callable[..., Any]:
|
|
1520
|
+
bursts: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _AsyncBurst] = (
|
|
1521
|
+
weakref.WeakKeyDictionary()
|
|
1522
|
+
)
|
|
1523
|
+
# Several loops (in several threads) share `bursts`; each burst itself is only ever
|
|
1524
|
+
# touched from its own loop's thread.
|
|
1525
|
+
bursts_lock = threading.Lock()
|
|
1526
|
+
|
|
1527
|
+
async def run(burst: _AsyncBurst) -> Any:
|
|
1528
|
+
loop = asyncio.get_running_loop()
|
|
1529
|
+
with bursts_lock:
|
|
1530
|
+
if bursts.get(loop) is burst:
|
|
1531
|
+
del bursts[loop]
|
|
1532
|
+
call = burst.call
|
|
1533
|
+
if call is None or burst.outcome.done():
|
|
1534
|
+
return None
|
|
1535
|
+
try:
|
|
1536
|
+
result = await func(*call[0], **call[1])
|
|
1537
|
+
except asyncio.CancelledError:
|
|
1538
|
+
burst.outcome.cancel()
|
|
1539
|
+
raise
|
|
1540
|
+
except BaseException as exc:
|
|
1541
|
+
if not burst.outcome.done():
|
|
1542
|
+
burst.outcome.set_exception(exc)
|
|
1543
|
+
burst.outcome.exception() # mark retrieved: callers may all be gone
|
|
1544
|
+
raise
|
|
1545
|
+
if not burst.outcome.done():
|
|
1546
|
+
burst.outcome.set_result(result)
|
|
1547
|
+
return result
|
|
1548
|
+
|
|
1549
|
+
async def sleep_then_run(burst: _AsyncBurst) -> None:
|
|
1550
|
+
await asyncio.sleep(wait) # real time, like the sync `threading.Timer`
|
|
1551
|
+
burst.timer = None # past the cancellation point: run even if called again now
|
|
1552
|
+
with contextlib.suppress(Exception): # delivered to callers via burst.outcome
|
|
1553
|
+
await run(burst)
|
|
1554
|
+
|
|
1555
|
+
@functools.wraps(func)
|
|
1556
|
+
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
1557
|
+
loop = asyncio.get_running_loop()
|
|
1558
|
+
with bursts_lock:
|
|
1559
|
+
burst = bursts.get(loop)
|
|
1560
|
+
if burst is None:
|
|
1561
|
+
burst = bursts[loop] = _AsyncBurst(loop)
|
|
1562
|
+
if burst.timer is not None:
|
|
1563
|
+
burst.timer.cancel()
|
|
1564
|
+
burst.call = (args, kwargs)
|
|
1565
|
+
burst.timer = loop.create_task(sleep_then_run(burst))
|
|
1566
|
+
return await asyncio.shield(burst.outcome)
|
|
1567
|
+
|
|
1568
|
+
def cancel_burst(burst: _AsyncBurst) -> None:
|
|
1569
|
+
if burst.timer is not None:
|
|
1570
|
+
burst.timer.cancel()
|
|
1571
|
+
burst.outcome.cancel()
|
|
1572
|
+
|
|
1573
|
+
def cancel() -> None:
|
|
1574
|
+
try:
|
|
1575
|
+
running: asyncio.AbstractEventLoop | None = asyncio.get_running_loop()
|
|
1576
|
+
except RuntimeError:
|
|
1577
|
+
running = None
|
|
1578
|
+
with bursts_lock:
|
|
1579
|
+
if running is not None:
|
|
1580
|
+
taken = [(running, bursts.pop(running, None))]
|
|
1581
|
+
else:
|
|
1582
|
+
taken = [(lp, bursts.pop(lp)) for lp in list(bursts)]
|
|
1583
|
+
for lp, burst in taken:
|
|
1584
|
+
if burst is None:
|
|
1585
|
+
continue
|
|
1586
|
+
if lp.is_closed():
|
|
1587
|
+
continue
|
|
1588
|
+
if lp is running or not lp.is_running():
|
|
1589
|
+
cancel_burst(burst)
|
|
1590
|
+
else:
|
|
1591
|
+
# Another thread's loop: futures are not thread-safe, so hand the
|
|
1592
|
+
# cancellation to that loop (which also wakes it up).
|
|
1593
|
+
with contextlib.suppress(RuntimeError): # closed meanwhile
|
|
1594
|
+
lp.call_soon_threadsafe(cancel_burst, burst)
|
|
1595
|
+
|
|
1596
|
+
async def flush() -> Any:
|
|
1597
|
+
with bursts_lock:
|
|
1598
|
+
burst = bursts.get(asyncio.get_running_loop())
|
|
1599
|
+
if burst is None:
|
|
1600
|
+
return None
|
|
1601
|
+
if burst.timer is not None:
|
|
1602
|
+
burst.timer.cancel()
|
|
1603
|
+
burst.timer = None
|
|
1604
|
+
return await run(burst)
|
|
1605
|
+
|
|
1606
|
+
wrapper.cancel = cancel # type: ignore[attr-defined]
|
|
1607
|
+
wrapper.flush = flush # type: ignore[attr-defined]
|
|
1608
|
+
return wrapper
|
|
1609
|
+
|
|
1610
|
+
|
|
1611
|
+
def debounce(wait: float) -> _DebounceDecorator:
|
|
1612
|
+
"""Postpone running until ``wait`` seconds have passed since the last call.
|
|
1613
|
+
|
|
1614
|
+
Each call restarts the countdown and replaces the pending arguments, so only the last
|
|
1615
|
+
call of a burst runs.
|
|
1616
|
+
|
|
1617
|
+
* Sync functions: the countdown is a daemon `threading.Timer`, so the function
|
|
1618
|
+
runs in the timer's thread. The wrapper returns ``None`` immediately.
|
|
1619
|
+
* ``async def`` functions: a call schedules a task on the running loop (cancelling the
|
|
1620
|
+
previously scheduled one) and then waits for the burst to finish. **Every caller in
|
|
1621
|
+
a burst awaits the same outcome** — the result (or exception) of the final call —
|
|
1622
|
+
so superseded callers are not left hanging or cancelled. Calls must be awaited (or
|
|
1623
|
+
wrapped in a task) to take effect; a single task awaiting sequentially simply waits
|
|
1624
|
+
``wait`` seconds per call. Bursts are tracked per event loop.
|
|
1625
|
+
|
|
1626
|
+
The wrapper also has:
|
|
1627
|
+
|
|
1628
|
+
* ``.cancel()`` — drop the pending call (async awaiters get ``CancelledError``).
|
|
1629
|
+
Thread-safe. For ``async def`` functions, inside a running loop it cancels that
|
|
1630
|
+
loop's pending call; from a thread with no running loop it cancels the pending
|
|
1631
|
+
call of every loop (handing the cancellation to each loop's own thread).
|
|
1632
|
+
* ``.flush()`` — run the pending call right now and return its result, or ``None``
|
|
1633
|
+
if nothing is pending. For ``async def`` functions it is a coroutine to ``await``.
|
|
1634
|
+
|
|
1635
|
+
State belongs to the decorated function, so on a method it is shared by all
|
|
1636
|
+
instances: calls on different instances debounce each other.
|
|
1637
|
+
|
|
1638
|
+
Args:
|
|
1639
|
+
wait: Quiet period in seconds; must be non-negative.
|
|
1640
|
+
|
|
1641
|
+
Returns:
|
|
1642
|
+
A decorator producing a `Debounced` or `AsyncDebounced` wrapper.
|
|
1643
|
+
|
|
1644
|
+
Raises:
|
|
1645
|
+
ValueError: If ``wait`` is negative.
|
|
1646
|
+
|
|
1647
|
+
Example:
|
|
1648
|
+
>>> @debounce(0.3)
|
|
1649
|
+
... def save(text: str) -> None: ...
|
|
1650
|
+
>>> save("a"); save("ab"); save.flush() # runs once, with "ab"
|
|
1651
|
+
"""
|
|
1652
|
+
if wait < 0:
|
|
1653
|
+
raise ValueError(f"debounce wait must be non-negative, got {wait!r}")
|
|
1654
|
+
|
|
1655
|
+
def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
|
|
1656
|
+
if _core.is_async(func):
|
|
1657
|
+
return _debounce_async(func, wait)
|
|
1658
|
+
return _debounce_sync(func, wait)
|
|
1659
|
+
|
|
1660
|
+
return cast(_DebounceDecorator, decorate)
|
|
1661
|
+
|
|
1662
|
+
|
|
1663
|
+
# ---------------------------------------------------------------------------------------
|
|
1664
|
+
# singleflight
|
|
1665
|
+
# ---------------------------------------------------------------------------------------
|
|
1666
|
+
|
|
1667
|
+
|
|
1668
|
+
class _KwargsMark:
|
|
1669
|
+
"""Private marker in keys built from keyword arguments (never equal to user data)."""
|
|
1670
|
+
|
|
1671
|
+
__slots__ = ()
|
|
1672
|
+
|
|
1673
|
+
def __repr__(self) -> str:
|
|
1674
|
+
return "<kwargs>"
|
|
1675
|
+
|
|
1676
|
+
|
|
1677
|
+
_KWARGS: Final = _KwargsMark()
|
|
1678
|
+
|
|
1679
|
+
|
|
1680
|
+
def _default_key(*args: Any, **kwargs: Any) -> Hashable:
|
|
1681
|
+
# The marker keeps ``f(1, b=2)`` and ``f((1,), (("b", 2),))`` apart.
|
|
1682
|
+
return (args, _KWARGS, tuple(sorted(kwargs.items()))) if kwargs else args
|
|
1683
|
+
|
|
1684
|
+
|
|
1685
|
+
class _Flights(Generic[T]):
|
|
1686
|
+
"""Thread-side bookkeeping for singleflight: key -> (in-flight future, leader thread)."""
|
|
1687
|
+
|
|
1688
|
+
def __init__(self) -> None:
|
|
1689
|
+
self.lock = threading.Lock()
|
|
1690
|
+
self.inflight: dict[Hashable, tuple[concurrent.futures.Future[T], int]] = {}
|
|
1691
|
+
|
|
1692
|
+
|
|
1693
|
+
class _AsyncFlight:
|
|
1694
|
+
"""One in-flight async singleflight run: its task and the callers waiting for it."""
|
|
1695
|
+
|
|
1696
|
+
__slots__ = ("task", "waiters")
|
|
1697
|
+
|
|
1698
|
+
def __init__(self) -> None:
|
|
1699
|
+
self.task: asyncio.Task[None] | None = None
|
|
1700
|
+
self.waiters: list[asyncio.Future[Any]] = []
|
|
1701
|
+
|
|
1702
|
+
def deliver(self, result: Any = None, exc: BaseException | None = None) -> None:
|
|
1703
|
+
"""Hand the outcome to every caller still waiting (``exc=None`` means success)."""
|
|
1704
|
+
waiters, self.waiters = self.waiters, []
|
|
1705
|
+
for waiter in waiters:
|
|
1706
|
+
if waiter.done():
|
|
1707
|
+
continue
|
|
1708
|
+
if exc is None:
|
|
1709
|
+
waiter.set_result(result)
|
|
1710
|
+
elif isinstance(exc, asyncio.CancelledError):
|
|
1711
|
+
waiter.cancel()
|
|
1712
|
+
else:
|
|
1713
|
+
waiter.set_exception(exc)
|
|
1714
|
+
|
|
1715
|
+
|
|
1716
|
+
def _start_task(
|
|
1717
|
+
loop: asyncio.AbstractEventLoop, coro: Coroutine[Any, Any, None]
|
|
1718
|
+
) -> asyncio.Task[None]:
|
|
1719
|
+
"""Start the shared run, eagerly where possible.
|
|
1720
|
+
|
|
1721
|
+
On Python 3.12+ with no custom task factory the task starts eagerly: it runs right
|
|
1722
|
+
away, up to its first suspension, so a run that never suspends (an in-memory hit,
|
|
1723
|
+
say) has finished before the caller would have waited at all. Otherwise (3.10/3.11,
|
|
1724
|
+
or a task factory is installed, which is respected) it starts on the next loop
|
|
1725
|
+
iteration, as `loop.create_task` does.
|
|
1726
|
+
"""
|
|
1727
|
+
if _EAGER_TASKS and loop.get_task_factory() is None:
|
|
1728
|
+
return asyncio.Task(coro, loop=loop, eager_start=True) # type: ignore[call-arg,unused-ignore]
|
|
1729
|
+
return loop.create_task(coro)
|
|
1730
|
+
|
|
1731
|
+
|
|
1732
|
+
_EAGER_TASKS = sys.version_info >= (3, 12)
|
|
1733
|
+
|
|
1734
|
+
|
|
1735
|
+
async def _await_flight(flight: _AsyncFlight, waiter: asyncio.Future[T]) -> T:
|
|
1736
|
+
"""Wait for a flight's outcome; a cancelled caller detaches, the run carries on.
|
|
1737
|
+
|
|
1738
|
+
The run resolves ``waiter`` itself (no done-callback relay), so an uncontended call
|
|
1739
|
+
costs one loop iteration for the run plus the wake-up, and none at all when the
|
|
1740
|
+
task ran eagerly to completion.
|
|
1741
|
+
"""
|
|
1742
|
+
try:
|
|
1743
|
+
return await waiter
|
|
1744
|
+
except asyncio.CancelledError:
|
|
1745
|
+
if waiter.done() and not waiter.cancelled():
|
|
1746
|
+
waiter.exception() # outcome arrived just as we were cancelled: retrieved
|
|
1747
|
+
else:
|
|
1748
|
+
with contextlib.suppress(ValueError):
|
|
1749
|
+
flight.waiters.remove(waiter)
|
|
1750
|
+
raise
|
|
1751
|
+
|
|
1752
|
+
|
|
1753
|
+
@overload
|
|
1754
|
+
def singleflight(func: F, /) -> F: ...
|
|
1755
|
+
@overload
|
|
1756
|
+
def singleflight(
|
|
1757
|
+
func: None = None, /, *, key: Callable[..., Hashable] | None = None
|
|
1758
|
+
) -> Callable[[F], F]: ...
|
|
1759
|
+
def singleflight(
|
|
1760
|
+
func: Callable[..., Any] | None = None,
|
|
1761
|
+
/,
|
|
1762
|
+
*,
|
|
1763
|
+
key: Callable[..., Hashable] | None = None,
|
|
1764
|
+
) -> Any:
|
|
1765
|
+
"""Let concurrent calls with the same arguments share a single execution.
|
|
1766
|
+
|
|
1767
|
+
While a call for a given key is running, further calls with that key do not start
|
|
1768
|
+
the function again; they wait and receive the same result, or the same exception.
|
|
1769
|
+
Once it finishes, the next call runs fresh — this de-duplicates work in flight, it
|
|
1770
|
+
is not a cache.
|
|
1771
|
+
|
|
1772
|
+
Sync functions de-duplicate across threads. ``async def`` functions de-duplicate
|
|
1773
|
+
across tasks on the same event loop: the work runs in a shared task, so cancelling
|
|
1774
|
+
one waiting caller does not cancel it for the others.
|
|
1775
|
+
|
|
1776
|
+
A re-entrant call — the function calling itself with the same key, on the leader's
|
|
1777
|
+
own thread (sync) or from inside the shared task (async) — cannot wait for its own
|
|
1778
|
+
result, so it runs the function directly, as if it were not decorated.
|
|
1779
|
+
|
|
1780
|
+
Args:
|
|
1781
|
+
func: The function to wrap (bare ``@singleflight`` form).
|
|
1782
|
+
key: Function of the call's arguments returning a hashable key. Defaults to the
|
|
1783
|
+
positional arguments plus sorted keyword arguments (including ``self`` for
|
|
1784
|
+
methods), which must then be hashable.
|
|
1785
|
+
|
|
1786
|
+
Returns:
|
|
1787
|
+
The wrapped function, or a decorator when called with options only.
|
|
1788
|
+
|
|
1789
|
+
Raises:
|
|
1790
|
+
TypeError: If the default key is used and an argument is unhashable.
|
|
1791
|
+
|
|
1792
|
+
Example:
|
|
1793
|
+
>>> @singleflight
|
|
1794
|
+
... def load_config(path: str) -> str:
|
|
1795
|
+
... return open(path).read()
|
|
1796
|
+
"""
|
|
1797
|
+
key_func = key or _default_key
|
|
1798
|
+
|
|
1799
|
+
def decorate(f: Callable[P, R]) -> Callable[P, R]:
|
|
1800
|
+
if _core.is_async(f):
|
|
1801
|
+
async_f = cast(Callable[P, Awaitable[Any]], f)
|
|
1802
|
+
flights_by_loop: weakref.WeakKeyDictionary[
|
|
1803
|
+
asyncio.AbstractEventLoop, dict[Hashable, _AsyncFlight]
|
|
1804
|
+
] = weakref.WeakKeyDictionary()
|
|
1805
|
+
loops_lock = threading.Lock() # loops in several threads share the mapping
|
|
1806
|
+
|
|
1807
|
+
@functools.wraps(f)
|
|
1808
|
+
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Any:
|
|
1809
|
+
k = key_func(*args, **kwargs)
|
|
1810
|
+
loop = asyncio.get_running_loop()
|
|
1811
|
+
inflight = flights_by_loop.get(loop)
|
|
1812
|
+
if inflight is None:
|
|
1813
|
+
with loops_lock:
|
|
1814
|
+
inflight = flights_by_loop.setdefault(loop, {})
|
|
1815
|
+
flight = inflight.get(k)
|
|
1816
|
+
if flight is not None and flight.task is asyncio.current_task():
|
|
1817
|
+
# Re-entrant call from inside the shared task: waiting for ourselves
|
|
1818
|
+
# would deadlock, so run the function directly.
|
|
1819
|
+
return await async_f(*args, **kwargs)
|
|
1820
|
+
waiter: asyncio.Future[Any] = loop.create_future()
|
|
1821
|
+
if flight is not None:
|
|
1822
|
+
flight.waiters.append(waiter)
|
|
1823
|
+
return await _await_flight(flight, waiter)
|
|
1824
|
+
|
|
1825
|
+
new = inflight[k] = _AsyncFlight()
|
|
1826
|
+
new.waiters.append(waiter) # before the task starts: it may run eagerly
|
|
1827
|
+
|
|
1828
|
+
async def run() -> None:
|
|
1829
|
+
new.task = asyncio.current_task()
|
|
1830
|
+
try:
|
|
1831
|
+
result = await async_f(*args, **kwargs)
|
|
1832
|
+
except BaseException as exc:
|
|
1833
|
+
if inflight.get(k) is new:
|
|
1834
|
+
del inflight[k]
|
|
1835
|
+
new.deliver(exc=exc)
|
|
1836
|
+
if not isinstance(exc, Exception):
|
|
1837
|
+
raise # CancelledError, KeyboardInterrupt, SystemExit
|
|
1838
|
+
# Delivered to the callers; the task itself ends quietly, so an
|
|
1839
|
+
# orphaned failure is never logged as "never retrieved".
|
|
1840
|
+
else:
|
|
1841
|
+
if inflight.get(k) is new:
|
|
1842
|
+
del inflight[k]
|
|
1843
|
+
new.deliver(result)
|
|
1844
|
+
|
|
1845
|
+
# Keep a strong reference: the event loop only holds tasks weakly.
|
|
1846
|
+
new.task = _start_task(loop, run())
|
|
1847
|
+
return await _await_flight(new, waiter)
|
|
1848
|
+
|
|
1849
|
+
return cast(Callable[P, R], async_wrapper)
|
|
1850
|
+
|
|
1851
|
+
flights: _Flights[R] = _Flights()
|
|
1852
|
+
|
|
1853
|
+
@functools.wraps(f)
|
|
1854
|
+
def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
|
|
1855
|
+
k = key_func(*args, **kwargs)
|
|
1856
|
+
me = threading.get_ident()
|
|
1857
|
+
with flights.lock:
|
|
1858
|
+
entry = flights.inflight.get(k)
|
|
1859
|
+
if entry is None:
|
|
1860
|
+
future: concurrent.futures.Future[R] = concurrent.futures.Future()
|
|
1861
|
+
flights.inflight[k] = (future, me)
|
|
1862
|
+
leader = True
|
|
1863
|
+
else:
|
|
1864
|
+
future, owner = entry
|
|
1865
|
+
leader = False
|
|
1866
|
+
if not leader:
|
|
1867
|
+
if owner == me:
|
|
1868
|
+
# Re-entrant call on the leader's own thread: waiting for our own
|
|
1869
|
+
# result would deadlock, so run the function directly.
|
|
1870
|
+
return f(*args, **kwargs)
|
|
1871
|
+
return future.result()
|
|
1872
|
+
try:
|
|
1873
|
+
result = f(*args, **kwargs)
|
|
1874
|
+
except BaseException as exc:
|
|
1875
|
+
with flights.lock:
|
|
1876
|
+
del flights.inflight[k]
|
|
1877
|
+
future.set_exception(exc)
|
|
1878
|
+
raise
|
|
1879
|
+
with flights.lock:
|
|
1880
|
+
del flights.inflight[k]
|
|
1881
|
+
future.set_result(result)
|
|
1882
|
+
return result
|
|
1883
|
+
|
|
1884
|
+
return wrapper
|
|
1885
|
+
|
|
1886
|
+
return decorate if func is None else decorate(func)
|