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.
@@ -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)