1# ruff: noqa: E722,PLR0913 # ty: ignore[unresolved-attribute]
2__lazy_modules__ = frozenset(('functools',))
3from asyncutils.config import _randinst
4from asyncutils.constants import _NO_DEFAULT
5from asyncutils._internal import log, patch as P
6from asyncutils._internal.helpers import fullname, get_loop_and_set
7from asyncutils._internal.submodules import func_all as __all__
8import asyncio as I, asyncutils as A
9from collections import deque, namedtuple
10from functools import partial, update_wrapper, wraps
11from itertools import count, repeat
12from sys import audit
13from time import perf_counter
[docs]
14def acompose(*F, wrap_last=True, _=I.iscoroutine):
15 async def g(*a, **k):
16 if _(r := next(i := reversed(F))(*a, **k)): r = await r
17 for f in i:
18 if _(r := f(r)): r = await r
19 return r
20 if wrap_last: update_wrapper(g, F[-1])
21 return g
[docs]
22async def areduce(f, it, initial=_NO_DEFAULT, *, await_=True):
23 it = A.iter_to_agen(it)
24 if initial is _NO_DEFAULT:
25 try: initial = await anext(it)
26 except StopAsyncIteration: raise TypeError('asyncutils.iters.areduce: empty (async) iterable passed without initial value') from None
27 if await_:
28 async for i in it: initial = await f(initial, i)
29 else:
30 async for i in it: initial = f(initial, i)
31 return initial
[docs]
32def star(f, /):
33 async def g(a=(), k=None, /): return await f(*a, **(k or {}))
34 return wraps(f)(g)
[docs]
35def unstar(f, /):
36 async def g(*a, **k): return await f(a, k)
37 return wraps(f)(g)
[docs]
38def every(interval, /, *, stop_when=None, count_f=True, verbose=False, stop_on_exc=True, wait_first=False, loop=None, max_iterations=None, timer=perf_counter, supplied_args=(), supplied_kwargs=None, default=_NO_DEFAULT, default_fname='<name unknown>', _='func.every: periodic coroutine %s reached the maximum of %d iterations'):
39 if loop is None: loop = get_loop_and_set()
40 def dec(f, /):
41 n = getattr(f, '__qualname__', default_fname)
42 if stop_when and stop_when.done(): log.warning('func.every: future to stop periodic coroutine %s is already done', n)
43 async def g(*a, **k):
44 log.debug('func.every: periodic task started'); q = default is _NO_DEFAULT; nonlocal stop_when
45 if stop_when is None: stop_when = loop.create_future()
46 if wait_first: await I.sleep(interval)
47 for i in count() if max_iterations is None else range(max_iterations):
48 t = timer()
49 try: await f(*supplied_args, *a, **(supplied_kwargs or {}), **k)
50 except A.CRITICAL: raise A.Critical
51 except:
52 if stop_on_exc:
53 if stop_when.done(): return stop_when.result()
54 break
55 (log.error if verbose else log.warning)('func.every: error in periodic coroutine %s on iteration %d', n, i, exc_info=True)
56 try: return await I.wait_for(stop_when, interval+t-timer() if count_f else interval)
57 except I.CancelledError:
58 if stop_on_exc: break
59 (log.info if verbose else log.debug)('func.every: future to stop periodic coroutine %s was cancelled on iteration %d', n, i, exc_info=True); stop_when = loop.create_future()
60 except TimeoutError: continue
61 else:
62 T = n, max_iterations
63 if stop_on_exc or default is A.RAISE: raise A.MaxIterationsError(_%T)
64 (log.info if verbose or q else log.debug)(_, *T)
65 if not q: return default
66 return wraps(f)(g)
67 return dec
[docs]
68def everymethod(interval, /, *, stop_when_getter=None, count_f=True, verbose=False, stop_on_exc=True, wait_first=False, loop=None, max_iterations=None, timer=perf_counter, supplied_args=(), supplied_kwargs=None, default=_NO_DEFAULT, default_fname='<name unknown>', _='func.everymethod: periodic coroutine %s reached the maximum of %d iterations'):
69 if loop is None: loop = get_loop_and_set()
70 def dec(f, /):
71 n = getattr(f, '__qualname__', default_fname)
72 async def g(self, /, *a, **k):
73 log.debug('func.everymethod: periodic task started'); q = default is _NO_DEFAULT
74 if (stop_when := loop.create_future() if stop_when_getter is None else stop_when_getter(self)).done(): log.warning('func.everymethod: future to stop periodic coroutine %s is already done', n)
75 if wait_first: await I.sleep(interval)
76 for i in count() if max_iterations is None else range(max_iterations):
77 t = timer()
78 try: await f(self, *supplied_args, *a, **(supplied_kwargs or {}), **k)
79 except A.CRITICAL: raise A.Critical
80 except:
81 if stop_on_exc:
82 if stop_when.done(): return stop_when.result()
83 break
84 (log.error if verbose else log.warning)('func.everymethod: error in periodic coroutine %s on iteration %d', n, i, exc_info=True)
85 try: return await I.wait_for(stop_when, interval+t-timer() if count_f else interval)
86 except I.CancelledError:
87 if stop_on_exc: break
88 (log.info if verbose else log.debug)('func.everymethod: future to stop periodic coroutine %s was cancelled on iteration %d', n, i, exc_info=True); stop_when = loop.create_future()
89 except TimeoutError: continue
90 else:
91 T = n, max_iterations
92 if stop_on_exc or default is A.RAISE: raise A.MaxIterationsError(_%T)
93 (log.info if verbose or q else log.debug)(_, *T)
94 if not q: return default
95 return wraps(f)(g)
96 return dec
[docs]
97def timer(f, /, *, precision=None, expected=Exception, should_log=True, timer=perf_counter, ns=False, _='nano', c='func.timer: function %s finished in %.*f %sseconds.', d='func.timer: received expected error from function %s after %.*f %sseconds: %s'): # cspell:disable-line
98 if precision is None: precision = A.getcontext().TIMER_DEFAULT_PRECISION-ns*9
99 async def g(*a, **k):
100 s = timer()
101 try:
102 r = await f(*a, **k); e = timer()-s
103 if should_log: log.info(c, fullname(f), precision, e, _ if ns else '')
104 return r, e
105 except A.CRITICAL: raise A.Critical
106 except expected as b:
107 e = timer()-s
108 if should_log: log.warning(d, fullname(f), precision, e, _ if ns else '', b, exc_info=True)
109 return A.wrap_exc(b), e
110 return wraps(f)(g)
[docs]
111def retry(tries=None, delay=None, *, max_delay=None, backoff=None, jitter=None, exc=Exception, on_retry=(_ := lambda *_: None), on_success=_, random=_randinst.random):
112 c = A.getcontext()
113 if tries is None: tries = c.RETRY_DEFAULT_TRIES
114 if delay is None: delay = c.RETRY_DEFAULT_DELAY
115 if backoff is None: backoff = c.RETRY_DEFAULT_BACKOFF
116 if max_delay is None: max_delay = c.RETRY_DEFAULT_MAX_DELAY
117 if jitter is None: jitter = c.RETRY_DEFAULT_JITTER
118 def dec(f):
119 async def g(*a, **k):
120 c, l, b = 0, 0.0, 1
121 for i in range(tries-1):
122 try: r = await f(*a, **k)
123 except exc as e:
124 c += 1
125 if I.iscoroutine(t := on_retry(i, e)): await t
126 await I.sleep(l := min(max(delay*b+(c*delay)*(1+(random()*2-1)*jitter), delay), max_delay)); b *= backoff
127 else:
128 if I.iscoroutine(t := on_success(i, l)): await t
129 return r
130 return await f(*a, **k)
131 return wraps(f)(g)
132 return dec
[docs]
133def throttle(lim, timer=perf_counter):
134 l = 0.0
135 def dec(f, /):
136 async def g(*a, **k):
137 nonlocal l
138 if w := max(0, 1/lim-timer()+l): await I.sleep(w)
139 l = timer(); return await f(*a, **k)
140 return wraps(f)(g)
141 return dec
[docs]
142def debounce(wait):
143 def dec(f, /, l=None):
144 (L := get_loop_and_set()).set_task_factory(I.eager_task_factory); g, h = L.create_task, I.sleep.__get__(wait)
145 async def j(*a, **k):
146 nonlocal l
147 if l: await A.safe_cancel(l)
148 l = g(h())
149 with A.ignore_cancellation: await l; return await f(*a, **k)
150 return wraps(f)(j)
151 return dec
[docs]
152def iterf(n, /):
153 def dec(f, /):
154 async def g(x, /):
155 for _ in repeat(None, n): x = await f(x)
156 return x
157 return wraps(f)(g)
158 return dec
[docs]
159async def measure(f, /, *, timer=perf_counter): s = timer(); return await f(), timer()-s
[docs]
160async def measure2(f, /, **k): return (await measure(f, **k))[1]
[docs]
161async def benchmark(f, /, times=None, warmup=None, _f=namedtuple('BenchmarkResult', 'min max total avg iterations', module='asyncutils.func'), *, sequential=None):
162 c, g = A.getcontext(), measure2.__get__(f)
163 if sequential is None: sequential = c.BENCHMARK_DEFAULT_SEQUENTIAL
164 if times is None: times = c.BENCHMARK_DEFAULT_TIMES
165 if warmup is None: warmup = c.BENCHMARK_DEFAULT_WARMUP
166 if sequential:
167 for _ in repeat(None, warmup): await f()
168 else: await I.gather(*(f() for _ in repeat(None, warmup)))
169 audit('asyncutils.func.benchmark', fullname(f), T := times+warmup); return _f(min(t := [await g() for _ in repeat(None, times)] if sequential else await I.gather(*(g() for _ in repeat(None, times)))), max(t), S := sum(t), S/times, T)
170P.patch_function_signatures((measure, _ := 'f, /, *, timer={}'), (measure, _), (benchmark, 'f, /, times=None, warmup=None'))
[docs]
171class RateLimited:
172 __slots__ = '__calls', '__ct', '__func', '__lock', '__period', '__raise', '__timer'
173 def __new__(cls, f, /, calls, period=None, *, raise_=False, timer=perf_counter, lock_impl=None):
174 if period is None: return partial(cls, calls=f, period=calls, raise_=raise_, timer=timer, lock_impl=lock_impl)
175 audit('asyncutils.func.RateLimited', fullname(f), calls, period); (_ := super().__new__(cls)).__func, _.__period, _.__ct, _.__lock, _.__calls, _.__raise, _.__timer = f, float(period), deque(), (I.Lock if lock_impl is None else lock_impl)(), int(calls), raise_, timer; return _
[docs]
176 async def __call__(self, *a, **k):
177 p, m, P, C, f = (T := self.__ct).popleft, T.appendleft, self.__period, self.__calls, self.__func
178 async with self.__lock:
179 d = (n := self.__timer())-P
180 while T:
181 if (x := p()) > d: m(x); break
182 if (l := len(T)-self.__calls+1) > 0:
183 if self.__raise: raise A.RateLimitExceeded(f, a, k, C, P, l)
184 await I.sleep(p()-d)
185 T.append(n)
186 return await f(*a, **k)
187 def __repr__(self): return f'{fullname(self)}({self.__func!r}, {self.__calls}, {self.__period:.6f}, raise_={self.__raise}, timer={self.__timer!r}, lock_impl={fullname(self.__lock)})'
188 P.patch_classmethod_signatures((__new__, 'f, /, calls, period=None, *, raise_=False, timer={}, lock_impl=None'))
189del _, perf_counter, P