1# ruff: noqa: RUF029
2from asyncutils import aenumerate, getcontext, iter_to_agen
3from asyncutils.config import _randinst
4from asyncutils.constants import _NO_DEFAULT
5from asyncutils._internal import py312 as Z, compat as C, helpers as H, patch as P
6from asyncutils._internal.submodules import iters_all as __all__
7import asyncutils as A, asyncio as B, operator as O, math as M
8from collections import Counter, defaultdict, deque, namedtuple
9from functools import partial, lru_cache, wraps
10from itertools import count, repeat
11from sys import audit, maxsize
12from time import monotonic
13_get0, _get1 = map(O.itemgetter, range(2))
14_rand, _randrange, _sample, _small_primes, _perfect_test, _identity = _randinst.random, _randinst.randrange, _randinst.sample, frozenset(_little_primes := (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199)), ((0x7ff, (2,)), (0x8a8d7f, (31, 73)), (0x11baa74c5, (2, 7, 61)), (0x1053cb094c1, (2, 13, 23, 0x195f53)), (0x1f51f3fee3b, _little_primes[:5]), (0x32907381cdf, _little_primes[:6]), (1<<64, (2, 0x145, 0x249f, 0x6e12, 0x6e0d7, 0x953d18, 0x6b0191fe)), (0x2be6951adc5b22410a5fd, _little_primes[:13]), (0x4c16c7697197146a6b8eb49518c5, _little_primes[:18])), lambda _, /: _
[docs]
15def fmap(fs, /, *a, **k): return agather(f(*a, **k) async for f in iter_to_agen(fs))
[docs]
16async def fmap_sequential(fs, /, *a, **k):
17 async for f in iter_to_agen(fs): yield await f(*a, **k)
[docs]
18async def fmap_parallel(fs, /, *a, **k):
19 t = H.get_loop_and_set().create_task
20 for r in await to_list(t(f(*a, **k)) async for f in iter_to_agen(fs)): yield await r
[docs]
21async def map_on_map(outer, inner, it, *, inner_await=False, outer_await=False):
22 async for _ in amap(inner, it, await_=inner_await): yield await to_tuple(amap(outer, _, await_=outer_await))
[docs]
23def aevery(it, n, *, skip_first=False): return aislice(it, skip_first, None, n)
[docs]
24def aevery_other(it, *, skip_first=False): return aevery(it, 2, skip_first=skip_first)
[docs]
25async def agather(it_of_its, return_exceptions=False): return await B.gather(*await to_list(it_of_its), return_exceptions=return_exceptions)
[docs]
26def aawgenf2agenf(f, /):
27 async def g(*a, **k):
28 async for _ in await f(*a, **k): yield _
29 return wraps(f)(g)
30async def _tee_helper(Q, i, /): await B.gather(*(q.put(i) for q in Q), return_exceptions=True)
[docs]
31def tee(it, n=2, *, maxqsize=None, put_exc=None, loop=None, _=_tee_helper):
32 if n <= 0: raise ValueError('asyncutils.iters.tee: n must be positive')
33 if n == 1: return iter_to_agen(it),
34 C = getcontext()
35 if loop is None: loop = H.get_loop_and_set()
36 if put_exc is None: put_exc = C.TEE_DEFAULT_PUT_EXC
37 if maxqsize is None: maxqsize = C.TEE_DEFAULT_MAX_QSIZE
38 Q = tuple(Z.Queue(maxqsize) for _ in repeat(None, n))
39 async def r(q):
40 try:
41 while True:
42 i = await q.get()
43 if put_exc and A.exception_occurred(i): raise A.unwrap_exc(i)
44 yield i
45 except Z.QueueShutDown:
46 nonlocal n; n -= 1
47 if n == 0: await A.safe_cancel(t)
48 async def f():
49 h = _.__get__(Q)
50 try: await agather(amap(h, it))
51 except A.CRITICAL: raise A.Critical
52 except BaseException as e:
53 if put_exc: await h(A.wrap_exc(e))
54 else:
55 for q in Q: q.shutdown(True)
56 raise
57 finally:
58 for q in Q: q.shutdown()
59 t = loop.create_task(f()); return tuple(map(r, Q))
[docs]
60async def adouble_starmap(f, it, /, await_=False):
61 it = iter_to_agen(it)
62 if await_:
63 async for _ in it: yield await f(**_)
64 else:
65 async for _ in it: yield f(**_)
[docs]
66async def astarmap_with_kwds(f, it, /, await_=False):
67 it = iter_to_agen(it)
68 if await_:
69 async for a, k in it: yield await f(*a, **k)
70 else:
71 async for a, k in it: yield f(*a, **k)
[docs]
72async def aloops(n, i=1024):
73 if n is None: n = maxsize
74 elif n <= 0: return
75 m, n = divmod(n, i)
76 for _ in repeat(None, m):
77 for _ in repeat(None, i): yield
78 await A.yield_to_event_loop
79 for _ in repeat(None, n): yield
80async def _aunzip_put(*_):
81 for q, i in zip(*_, strict=True):
82 with A.ignore_qshutdown: await q.put(i)
[docs]
83async def aunzip(ait, *, fillvalue=_NO_DEFAULT, put_batch=None, maxqsize=None, _a=_aunzip_put, _b=_identity):
84 audit('asyncutils.iters.aunzip', H.fullname(ait)); l = len(t := await anext(ait := iter_to_agen(ait), ())); C = getcontext()
85 if maxqsize is None: maxqsize = C.AUNZIP_DEFAULT_MAX_QSIZE
86 if put_batch is None: put_batch = C.AUNZIP_DEFAULT_PUT_BATCH
87 if maxqsize < put_batch: raise ValueError('asyncutils.iters.aunzip: maxqsize cannot be less than put_batch')
88 f = partial(Z.Queue, maxqsize)
89 class AUnzipConsumer:
90 __slots__ = '__q',
91 def __init__(self): self.__q = f()
92 async def __anext__(self, l=B.Lock(), f=partial(A.take, ait, put_batch, default=A.RAISE)): # noqa: B008
93 if self.__q.empty():
94 async with l:
95 try:
96 async for _ in f(): await _a(Q, _)
97 except A.ItemsExhausted:
98 for q in Q: q.close()
99 try: r = await self.__q.get()
100 except Z.QueueShutDown: raise StopAsyncIteration from None
101 if r is fillvalue: raise StopAsyncIteration
102 return r
103 def close(self): self.__q.shutdown()
104 __aiter__, __anext__.__text_signature__ = _b, '($self)' # ty: ignore[unresolved-attribute]
105 await _a(Q := await to_tuple(AUnzipConsumer() async for _ in aloops(l)), t); return Q
[docs]
106async def merge(*I, reverse=False, maxqsize=None, _=lambda p: lambda i: aconsume(amap(p, i, await_=True))):
107 audit('asyncutils.iters.merge', I); p, g, l, a = (q := (Z.LifoQueue if reverse else Z.Queue)(getcontext().MERGE_DEFAULT_MAX_QSIZE if maxqsize is None else maxqsize)).put, q.get, H.get_loop_and_set(), object()
108 async def close():
109 await B.gather(*map(_(p), I))
110 if not reverse: await p(a)
111 if reverse: q.put_nowait(a)
112 t = l.create_task(close())
113 if reverse: await t
114 while True:
115 if (c := await g()) is a: break
116 yield c
[docs]
117def aflatten(it, _=A.AChain.from_iterable): return aiter(_(it))
[docs]
118def acountdown(n, step=1, *, include_zero=False): return arange(n, -include_zero, -step)
119async def _traverse(s, n, q, f, i, /):
120 a, c, g = (v := {s}).add, v.__contains__, q.append
121 if i: yield s
122 while q:
123 async for _ in afilterfalse(c, n(f())): a(_); g(_); yield _
[docs]
124def abfs(start, neighbours, *, _=_traverse, include_start=True): return _(start, lambda x, /, _=neighbours: areversed(_(x)), q := deque((start,)), q.popleft, include_start)
[docs]
125def adfs(start, neighbours, *, _=_traverse, include_start=True): return _(start, neighbours, q := [start], q.pop, include_start)
[docs]
126async def asattolo(it, /, _=_randrange):
127 i = len(a := await to_list(it))
128 while i > 1:
129 i -= 1
130 a[j], a[i] = a[i], a[j := _(i)]
131 return a
[docs]
132async def abrent(f, s, /):
133 p = l = 1; t, h, m = s, await f(s), 0
134 while t is not h:
135 if p == l: t, l, p = h, 0, p<<1
136 h, l = await f(h), l+1
137 a = s, await A.iterf(l)(f)(s)
138 while a[0] is not a[1]: a, m = await B.gather(*map(f, a)), m+1
139 return a[0], l, m
[docs]
140async def asample_l(it, k, *, rrange=_randrange, rand=_rand):
141 if k < 0: raise ValueError('asyncutils.iters.asample_l: expected non-negative sample size')
142 if k == 0: return []
143 R, W, i = await A.collect(it := iter_to_agen(it), k, A.RAISE), 1.0, k
144 while True:
145 W *= rand()**(1.0/k); i += (s := M.floor(M.log(rand(), 1-W))+1)
146 try: R[rrange(k)] = await anth(it, s)
147 except A.ItemsExhausted: return R
[docs]
148async def asample_weighted(it, k, *, rrange=_randrange, rand=_rand):
149 if k < 0: raise ValueError('asyncutils.iters.asample_weighted: expected non-negative sample size')
150 if k == 0: return []
151 W, u, p = 0.0, rand(), 1.0
152 async def agen(it):
153 nonlocal W
154 async for i, w in iter_to_agen(it):
155 if w < 0: raise ValueError(f'asyncutils.iters.asample_weighted: weight {w} for item {i!r} is negative')
156 W += w; yield i, w
157 r = await A.collect(it := agen(it), k, A.RAISE)
158 async for i, w in it:
159 w /= W; u -= w*p; p *= 1-w # noqa: PLW2901
160 if u <= 0: r[rrange(k)], u, p = i, rand(), 1.0
161 return r
[docs]
162async def astarfilter(pred, it, await_pred=False):
163 it = iter_to_agen(it)
164 if await_pred:
165 if pred is None: raise ValueError('asyncutils.iters.astarfilter: pred cannot be None if await_pred is True')
166 async for i in it:
167 if await pred(*await to_list(i)): yield i
168 else:
169 if pred is None: pred = bool
170 async for i in it:
171 if pred(*await to_list(i)): yield i
[docs]
172async def astarfilterfalse(pred, it, await_pred=False):
173 it = iter_to_agen(it)
174 if await_pred:
175 if pred is None: raise ValueError('asyncutils.iters.astarfilterfalse: pred cannot be None if await_pred is True')
176 async for i in it:
177 if not await pred(*await to_list(i)): yield i
178 else:
179 if pred is None: pred = bool
180 async for i in it:
181 if not pred(*await to_list(i)): yield i
[docs]
182def amultistarfilter(p, /, *i, strict=False, await_pred=False): return astarfilter(p, azip(*i, strict=strict), await_pred)
[docs]
183async def amultistarfilterfalse(p, /, *i, strict=False, await_pred=False): return astarfilterfalse(p, azip(*i, strict=strict), await_pred)
[docs]
184def amultifilter(p, /, *i, strict=False, await_pred=False): return afilter(p, azip(*i, strict=strict), await_pred)
[docs]
185def amultifilterfalse(p, /, *i, strict=False, await_pred=False): return afilterfalse(p, azip(*i, strict=strict), await_pred)
[docs]
186def hamming_dist(u, v, cmpeq=H.check): return ailen(amultistarfilterfalse(cmpeq, u, v))
[docs]
187async def amerge_sorted_by(its, *, key=None, await_key=False, reverse=False, _=A.ignore_stop_async_iteration):
188 f, a, b, c = (h := []).append, (m := C if reverse else __import__('heapq')).heapify, m.heappop, m.heappush
189 for i, it in enumerate(its := await to_tuple(amap(partial(iterate_with_key, key=key, await_key=await_key), its))):
190 with _: k, v = await anext(it); f((k, i, v))
191 a(h)
192 while h:
193 k, i, v = b(h); yield v
194 with _: k, v = await anext(its[i]); c(h, (k, i, v))
[docs]
195async def batch(it, n, *, item_timeout=None, strict=False):
196 f, g, _ = iter_to_agen(it).__anext__, (b := []).append, 0
197 while True:
198 for _ in range(n):
199 try: g(await B.wait_for(f(), item_timeout))
200 except StopAsyncIteration: break
201 except TimeoutError:
202 if b: break
203 if b:
204 if strict and _ < n-1: raise ValueError('asyncutils.iters.batch: incomplete batch')
205 yield H.copy_and_clear(b)
[docs]
206def batch2(it, n, strict=False): return A.aiter_from_f(partial(A.collect, it, n, default=A.RAISE if strict else _NO_DEFAULT), [])
[docs]
207async def aside_effect(f, it, size=None, *, await_=True):
208 if size is None:
209 it = iter_to_agen(it)
210 if await_:
211 async for i in it: await f(i); yield i
212 else:
213 async for i in it: f(i); yield i
214 else:
215 it = batch(it, size)
216 if await_:
217 async for i in it:
218 await f(i)
219 for _ in i: yield _
220 else:
221 async for i in it:
222 f(i)
223 for _ in i: yield _
[docs]
224def asliced(seq, n, strict=False):
225 I = atakewhile(None, (seq[i:i+n] async for i in acount(step=n)))
226 if not strict: return I
227 async def r():
228 async for s in I:
229 if len(s) != n: raise ValueError(f'asyncutils.iters.asliced: length of {seq!r} is not divisible by {n}')
230 yield s
231 return r()
232async def _buffer_consume(g, d, f, t, c, /):
233 x = t+c()
234 while True:
235 yield await g(); d()
236 if c() > x: await f(); x = t+c()
[docs]
237def buffer(it, maxsize=0, *, timeout_get=None, timeout_put=None, cooldown=0.0, _=_buffer_consume):
238 q = Z.Queue(maxsize)
239 async def cons():
240 try:
241 async for i in _(q.get, q.task_done, B.sleep.__get__(cooldown), float('inf') if timeout_get is None else timeout_get, monotonic): yield i
242 finally: await A.safe_cancel(t)
243 async def prod(p=q.put):
244 try:
245 async for _ in iter_to_agen(it):
246 try: await B.wait_for(p(_), timeout_put)
247 except TimeoutError: break
248 finally: c.close()
249 c, t = cons(), H.get_loop_and_set().create_task(prod()); return c
[docs]
250async def asplitat(it, pred, maxsplit=-1, keep_sep=False):
251 I, f = iter_to_agen(it), (b := []).append
252 if not maxsplit: yield await to_list(I); return
253 async for i in I:
254 if not pred(i): f(i); continue
255 yield b
256 if keep_sep: yield [i]
257 if maxsplit == 1: yield await to_list(I); return
258 f = (b := []).append; maxsplit -= 1
259 yield b
[docs]
260def batch_process(items, size, processor): return amap(processor, batch(items, size), await_=True)
[docs]
261async def window(it, size, step=1):
262 if not size > 0 < step: raise ValueError('asyncutils.iters.window: size and step should both be positive')
263 a, c = (b := deque(maxlen=size)).append, 0
264 async for i in iter_to_agen(it):
265 a(i)
266 if len(b) == size:
267 if not c%step and (t := (yield tuple(b))) is not None: size, step = t
268 c += 1
[docs]
269async def aall(it):
270 async for _ in iter_to_agen(it):
271 if not _: return False
272 return True
[docs]
273async def aany(it):
274 async for _ in iter_to_agen(it):
275 if _: return True
276 return False
[docs]
277async def aisempty(it):
278 async for _ in iter_to_agen(it): return False
279 return True
280async def _extreme(I, K, a, c, d, /):
281 try:
282 k, r = await anext(I := iterate_with_key(I[0] if len(I) == 1 else I, K, a))
283 except StopAsyncIteration:
284 if d is _NO_DEFAULT: raise ValueError('empty (async) iterable passed to asyncutils.iters.amax or asyncutils.iters.amin with no default value')
285 return d
286 async for i, x in I:
287 if c(x, k): k, r = x, i
288 return r
[docs]
289def amax(*i, key=None, default=_NO_DEFAULT, await_key=False, _=_extreme): return _(i, key, await_key, O.gt, default)
[docs]
290def amin(*i, key=None, default=_NO_DEFAULT, await_key=False, _=_extreme): return _(i, key, await_key, O.lt, default)
[docs]
291async def azip(*I, strict=False, _=A.ignore_stop_async_iteration.combined(RuntimeError)): # noqa: B008
292 if not I: return
293 I = tuple(map(iter_to_agen, I))
294 with _:
295 while True: yield tuple(await B.gather(*map(anext, I))) # noqa: ASYNC119
296 if not strict: return
297 for x, y in enumerate(I):
298 with _: await anext(y); raise ValueError(f'asyncutils.iters.azip: iterable {x} longer than shortest iterable')
[docs]
299async def amap(f, /, *i, await_=False, strict=False):
300 it = azip(*i, strict=strict)
301 if await_:
302 async for _ in it: yield await f(*_)
303 else:
304 async for _ in it: yield f(*_)
[docs]
305async def afilter(f, it, await_=False):
306 async for k, i in iterate_with_key(it, f, await_):
307 if k: yield i
[docs]
308def amapif(f, p, it, await_transform=False, await_pred=False): return amap(f, afilter(p, it, await_pred), await_=await_transform)
[docs]
309def amultimapif(f, /, *a, await_transform=False, await_pred=False): return astarmap(f, amultifilter(*a, await_pred=await_pred), await_transform)
[docs]
310def amultistarmapif(f, /, *a, await_transform=False, await_pred=False): return astarmap(f, amultistarfilter(*a, await_pred=await_pred), await_transform)
[docs]
311def arange(*a): return iter_to_agen(range(*a))
[docs]
312def acount(start=0, step=1):
313 if isinstance(step, float):
314 if step.is_integer(): step = int(step)
315 else: start = float(start)
316 elif start.is_integer(): start = int(start)
317 return iter_to_agen(count(start, step))
[docs]
318async def acycle(it):
319 if type(it) in C.s: l = tuple(it)
320 else:
321 a = (l := []).append
322 async for i in iter_to_agen(it): yield i; a(i)
323 l = tuple(l)
324 del it
325 while True:
326 for i in l: yield i
327 await A.yield_to_event_loop
[docs]
328async def arepeat(elem, n=None):
329 if n is None or n < 0:
330 while True: yield elem
331 else:
332 while n: yield elem; n -= 1
[docs]
333async def aaccumulate(it, func=O.add, *, initial=None):
334 it = iter_to_agen(it)
335 if initial is None:
336 try: initial = await anext(it)
337 except StopAsyncIteration: return
338 yield initial
339 async for _ in it: yield (initial := func(initial, _))
[docs]
340async def acompress(data, selectors):
341 async for i, j in azip(data, selectors):
342 if j: yield i
[docs]
343async def adropwhile(pred, it, await_pred=False):
344 async for p, _ in iterate_with_key(it := iter_to_agen(it), pred, await_pred):
345 if not p: yield _; break
346 async for _ in it: yield _
[docs]
347async def adropwhile_exclusive(pred, it, await_pred=False):
348 async for p, _ in iterate_with_key(it := iter_to_agen(it), pred, await_pred):
349 if not p: break
350 async for _ in it: yield _
[docs]
351async def adropuntil(pred, it, await_pred=False):
352 async for p, _ in iterate_with_key(it := iter_to_agen(it), pred, await_pred):
353 if p: yield _; break
354 async for _ in it: yield _
[docs]
355async def adropuntil_exclusive(pred, it, await_pred=False):
356 async for p, _ in iterate_with_key(it := iter_to_agen(it), pred, await_pred):
357 if p: break
358 async for _ in it: yield _
[docs]
359async def ac3merge(seqs):
360 seqs, g, d, c = await to_list(afilter(None, seqs)), (n := []).append, n.clear, None
361 while seqs:
362 for s in seqs:
363 c = s[0]
364 for t in seqs:
365 next(t := iter(t))
366 if any(H.check(_, c) for _ in t): c = None; break
367 else: break
368 if c is None: raise ValueError('asyncutils.iters.ac3merge: cannot resolve sequences')
369 yield c
370 for s in seqs:
371 if s[0] == c: del s[0]
372 if s: g(s)
373 seqs = tuple(n); d(); await A.yield_to_event_loop
[docs]
374async def afilterfalse(f, it, await_=False):
375 async for k, i in iterate_with_key(it, f, await_):
376 if not k: yield i
[docs]
377async def agroupby(it, key=None, await_=False):
378 it, c = iterate_with_key(it, key, await_), True
379 async def f(t):
380 nonlocal c, k, v; yield v
381 async for k, v in it:
382 if k != t: return
383 yield v
384 c = False
385 try: k, v = await anext(it)
386 except StopAsyncIteration: return
387 while c:
388 yield k, (g := f(t := k))
389 if k == t: await aconsume(g)
[docs]
390async def aislice(it, /, *a, _=lambda x: x if x is None else int(x, 0) if isinstance(x, str) else int(x)):
391 x, y, z = 0 if (s := slice(*map(_, a))).start is None else s.start, s.stop, 1 if s.step is None else s.step
392 if x < 0: raise ValueError(f'asyncutils.iters.aislice: start={x} is invalid')
393 if y is not None and y < 0: raise ValueError(f'asyncutils.iters.aislice: stop={y} is invalid')
394 if z <= 0: raise ValueError(f'asyncutils.iters.aislice: step={z} is invalid')
395 async for i, j in azip(acount() if y is None else arange(max(x, y)), it):
396 if i == x: yield j; x += z
[docs]
397async def aiter_idx(it, value, start=0, stop=None, _=H.check):
398 async for i, j in aenumerate(aislice(it, start, stop), start):
399 if _(j, value): yield i
[docs]
400async def asieve(n):
401 if n < 2: return # noqa: PLR2004
402 yield 2; s, d = 3, bytearray((0, 1))*(n>>1)
403 async for p in aiter_idx(d, 1, s, M.isqrt(n)+1):
404 async for i in aiter_idx(d, 1, s, s := p*p): yield i
405 d[s:n:x] = bytes(len(range(s, n, x := p<<1)))
406 async for i in aiter_idx(d, 1, s): yield i
[docs]
407async def apairwise(it):
408 try: a = await anext(I := iter_to_agen(it))
409 except StopAsyncIteration: return
410 async for b in I: yield a, b; a = b
[docs]
411@aawgenf2agenf
412async def atriplewise(it):
413 a, b, c = tee(iter_to_agen(it), 3, maxqsize=3); await B.gather(*(anext(g, None) for g in (b, c, c)))
414 return azip(a, b, c)
[docs]
415async def aproduct(*i, repeat=1):
416 if repeat < 0: raise ValueError('asyncutils.iters.aproduct: repeat cannot be negative')
417 r = [()]
418 async for p in arepeat(amap(to_tuple, i, await_=True), repeat): r = [(*x, y) for x in r async for y in p]
419 for _ in r: yield _
[docs]
420async def astarmap(f, it, /, await_=False):
421 it = iter_to_agen(it)
422 if await_:
423 async for _ in it: yield await f(*_)
424 else:
425 async for _ in it: yield f(*_)
[docs]
426async def atakewhile(pred, it, await_pred=False):
427 async for k, i in iterate_with_key(it, pred, await_pred):
428 if not k: break
429 yield i
[docs]
430async def atakewhile_inclusive(pred, it, await_pred=False):
431 async for k, i in iterate_with_key(it, pred, await_pred):
432 yield i
433 if not k: break
[docs]
434async def atakeuntil(pred, it, await_pred=False):
435 async for k, i in iterate_with_key(it, pred, await_pred):
436 if k: break
437 yield i
[docs]
438async def atakeuntil_inclusive(pred, it, await_pred=False):
439 async for k, i in iterate_with_key(it, pred, await_pred):
440 yield i
441 if k: break
[docs]
442def asum_of_squares(it): return asumprod(*tee(it))
[docs]
443async def aziplongest(*i, fillvalue=None):
444 n = len(i := list(map(iter_to_agen, i)))
445 while True:
446 f = (v := []).append
447 for j, a in enumerate(i):
448 try: x = await anext(a)
449 except StopAsyncIteration:
450 n -= 1
451 if not n: return
452 i[j], x = arepeat(fillvalue), fillvalue
453 f(x)
454 yield tuple(v)
[docs]
455def asumprod(p, q, /): return asum(amap(O.mul, p, q, strict=True))
[docs]
456async def aconvolve(signal, kernel, _=A.AChain):
457 f = (w := deque((0,), n := len(K := await to_tuple(areversed(kernel))))*n).append
458 async for x in _(signal, arepeat(0, n-1)): f(x); yield await asumprod(K, w)
[docs]
459def atabulate(f, start=0, step=1, /, *, await_=True): return amap(f, acount(start, step), await_=await_)
[docs]
460async def asum(it, start=0):
461 async for i in iter_to_agen(it): start += i
462 return start
[docs]
463async def aprod(it, start=1):
464 async for i in iter_to_agen(it): start *= i
465 return start
[docs]
466async def amatprod(it, start):
467 async for i in iter_to_agen(it): start @= i
468 return start
[docs]
469def atail(n, it, /): return aislice(it, max(0, len(it)-n), None)
[docs]
470async def to_tuple(it): return tuple(await to_list(it))
[docs]
471async def to_set(it, frozen=False): r = set(it) if type(it) in C.s else {_ async for _ in iter_to_agen(it)}; return frozenset(r) if frozen else r
[docs]
472async def to_list(it): return list(it) if type(it) in C.s else [_ async for _ in iter_to_agen(it)]
[docs]
473async def to_deque(it):
474 if type(it) in C.s: return deque(it)
475 a = (d := deque()).append
476 async for i in iter_to_agen(it): a(i)
477 return d
[docs]
478async def aconsume(it, n=None, _=H.check_methods):
479 if n == 0: return
480 if n: it = A.take(it, n, A.RAISE)
481 if _(it, '__iter__'): await H.get_loop_and_set().run_in_executor(H.create_executor(aconsume) if (E := getattr(aconsume, 'executor', None)) is None else E, deque, it, 0)
482 else:
483 async for _ in it: ...
[docs]
484def anth(it, n, default=_NO_DEFAULT): return anext(aislice(it, n, None), *H.filter_out(default, s=_NO_DEFAULT))
[docs]
485async def aall_equal(it, key=None, strict=False):
486 async for _ in (it := agroupby(it, key)):
487 async for _ in it: return False
488 return True
489 if strict: raise ValueError('asyncutils.aall_equal: iterable cannot be empty with strict=True')
490 return True
[docs]
491async def acombinations(it, r):
492 if r > (n := len(p := await to_tuple(it))): return
493 I = list(range(r)); yield p[:r]
494 while True:
495 for i in range(r-1, -1, -1):
496 if I[i] != i+n-r: break
497 else: return
498 I[i] += 1
499 for j in range(i+1, r): I[j] = I[j-1]+1
500 yield tuple(p[i] for i in I)
[docs]
501async def acombinations_with_replacement(it, r):
502 if not (n := len(p := await to_tuple(it))) and r: return
503 I = [0]*r; yield (p[0],)*r
504 while True:
505 for i in range(r-1, -1, -1):
506 if I[i] != n-1: break
507 else: return
508 I[i:] = (I[i]+1,)*(r-i); yield tuple(p[i] for i in I)
[docs]
509async def apermutations(it, r=None):
510 n = len(p := await to_tuple(it))
511 if (r := n if r is None else r) > n: return
512 I, C, x = list(range(n)), list(range(n, n-r, -1)), r-1; yield p[:r]
513 while n:
514 for i in range(x, -1, -1):
515 C[i] -= 1
516 if C[i]: I[i], I[-j] = I[-(j := C[i])], I[i]; yield tuple(p[i] for i in I[:r]); break
517 else: I[i:], C[i] = I[i+1:]+I[i:i+1], n-i
518 else: return
[docs]
519@aawgenf2agenf
520async def apowerset(it): s = await to_tuple(it); return aflatten(acombinations(s, r) for r in range(len(s)+1))
[docs]
521def aquantify(it, pred=bool): return asum(amap(pred, it))
[docs]
522async def apadded(it, fillvalue, n=None):
523 if n is None:
524 async for i in iter_to_agen(it): yield i
525 while True: yield fillvalue
526 async for n, j in A.aenumerate(it, n-1, step=-1): yield j # noqa: B007,B020,PLR1704
527 async for _ in aloops(n): yield fillvalue
[docs]
528def apadnone(it, n=None): return apadded(it, None, n)
[docs]
529def agrouper(it, n, fillvalue=_NO_DEFAULT): I = (iter_to_agen(it),)*n; return azip(*I, strict=fillvalue is A.RAISE) if isinstance(fillvalue, type(A.RAISE)) else aziplongest(*I, fillvalue=fillvalue)
[docs]
530async def aroundrobin(*i):
531 I = (iter_to_agen(_) for _ in i)
532 for j in range(len(i), 0, -1):
533 async for _ in (I := acycle(aislice(I, j))): yield await anext(_)
[docs]
534def aroundrobin2(*i): return afilter(partial(O.is_not, _NO_DEFAULT), aflatten(aziplongest(*i, fillvalue=_NO_DEFAULT)))
[docs]
535async def aunique_everseen(it, key=None, await_key=False):
536 A, a = (S := set()).add, (s := []).append
537 async for k, i in iterate_with_key(it, key, await_key):
538 try:
539 if k not in S: A(k); yield i
540 except TypeError:
541 if k not in s: a(k); yield i
[docs]
542def aunique_justseen(it, key=None): return amap(_get0, agroupby(it)) if key is None else amap(anext, amap(_get1, agroupby(it, key)), await_=True)
[docs]
543@aawgenf2agenf
544async def aunique(it, key=None, reverse=False): return aunique_justseen(await asorted(it, key=key, reverse=reverse), key)
[docs]
545@aawgenf2agenf
546async def ancycles(it, n): return aflatten(arepeat(await to_tuple(it), n))
[docs]
547def apartition(pred, it):
548 if pred is None: pred = bool
549 async def agen(q, _=iter_to_agen(it).__anext__): # noqa: B008
550 p = q.popleft
551 while True:
552 while q: yield p()
553 try: (T if pred(v := await _()) else F).append(v)
554 except StopAsyncIteration: return
555 return map(agen, (F := deque(), T := deque()))
[docs]
556async def aiterexcept(f, exc, first=None):
557 if first is not None: yield await first()
558 with A.IgnoreErrors(exc):
559 while True: yield await f() # noqa: ASYNC119
[docs]
560async def ailen(it):
561 i = 0
562 async for _ in iter_to_agen(it): i += 1
563 return i
[docs]
564async def aiterate(f, start):
565 while True: yield start; start = await f(start)
[docs]
566async def asorted(it, *, key=_identity, reverse=False, await_key=False):
567 g, f, b, a = (m := C if reverse else __import__('heapq')).heappop, m.heapify, (m := []).append, (r := []).append
568 if await_key:
569 async for i, j in aenumerate(it): b((await key(j), i, j))
570 else:
571 async for i, j in aenumerate(it): b((key(j), i, j))
572 f(m)
573 while m: a(g(m))
574 return r
[docs]
575def acanonical(it): return asorted(it, key=id, reverse=True)
576async def _pf(a, _):
577 while True:
578 yield tuple(a)
579 async for i in arange(_-2, -1, -1):
580 if a[i] < a[i+1]: break
581 else: return
582 async for j in arange(j := _-1, i, -1): # noqa: B020
583 if a[i] < a[j]: break
584 a[i], a[j] = a[j], a[i]; a[i+1:] = a[:i-_:-1]
585async def _pp(a, _):
586 h, R, l = a[:_], range(_-1, -1, -1), range(len(t := a[_:]))
587 while True:
588 yield tuple(h); p = t[-1]
589 for i in R:
590 if h[i] < p: break
591 p = h[i]
592 else: return
593 p = h[i]
594 for j in l:
595 if (c := t[j]) > p: h[i], t[j] = c, p; break
596 else:
597 for j in R:
598 if (c := h[j]) > p: h[i], h[j] = c, p; break
599 t += h[:-(x := _-i):-1]; i += 1; h[i:], t[:] = t[:x], t[x:]; await A.yield_to_event_loop
[docs]
600async def empty_agen(): return; yield
[docs]
601async def agives(x, /): yield x
[docs]
602@aawgenf2agenf
603async def adistinct_permutations(it, r=None, f=(_pp, _pf)):
604 if (S := len(I := await to_list(it))) < (_ := S if r is None else r): return agives(())
605 if _ <= 0: return empty_agen()
606 a = f[_ == S]
607 try: I.sort(); return a(I, _)
608 except TypeError:
609 d = defaultdict(list)
610 for i in I: d[I.index(i)].append(i)
611 return amap(lambda i, E={k: acycle(v) for k, v in d.items()}: to_tuple(await anext(E[_]) for _ in i), a(await asorted(amap(I.index, I)), _), await_=True) # noqa: B008
[docs]
612async def aunique_to_each(*i):
613 p = frozenset(await to_list(amap(to_tuple, i, await_=True)))
614 for x, j in Counter(await to_list(aflatten(map(frozenset, p)))).items():
615 if j == 1 and x in p: yield x
[docs]
616async def aderangements(it, r=None):
617 async for _ in acompress(apermutations(X := await to_tuple(it), r), amap(aall, amap(partial(amap, O.is_not), arepeat(Y := tuple(range(len(X)))), apermutations(Y, r)), await_=True)): yield _
[docs]
618def aintersperse(e, it, n=1):
619 if n <= 0: raise ValueError('asyncutils.iters.aintersperse: n must be positive')
620 return aislice(ainterleave_stopearly(arepeat(e), it), 1, None) if n == 1 else aflatten(aislice(ainterleave_stopearly(arepeat((e,)), batch(it, n)), 1, None))
[docs]
621def ainterleave_stopearly(*i): return aflatten(azip(*i))
[docs]
622def aspy(it, n=1): p, q = tee(it, maxqsize=n); return A.take(q, n), p
[docs]
623async def ainterleave_evenly(its, lengths=None):
624 I = await to_tuple(its)
625 try:
626 if (X := len(I)) != len(L := await to_tuple(lengths or map(len, I))): raise ValueError('asyncutils.iters.ainterleave_evenly: mismatch in length of its and lengths')
627 except TypeError: raise ValueError('asyncutils.iters.ainterleave_evenly: cannot determine lengths of (async) iterables') from None
628 A, *a = map(f := L.__getitem__, _ := sorted(range(X), key=f, reverse=True)); B, *b = (iter_to_agen(I[i]) for i in _); E, t = [A//X]*len(a), sum(L)
629 while t:
630 yield await anext(B); t -= 1; E[:] = map(O.sub, E, a)
631 for i, e in enumerate(E):
632 if e < 0: yield await anext(b[i]); t -= 1; E[i] += A
[docs]
633async def ainterleave_randomly(its, _=_randrange):
634 x = len(I := await to_list(amap(iter_to_agen, its)))
635 while x:
636 i = _(x)
637 try: yield await anext(I[i])
638 except StopAsyncIteration: I[i] = I[-1]; del I[-1]; x -= 1
[docs]
639async def acollapse(it, base_typ=(str, bytes), levels=None):
640 if levels is None: levels = float('inf')
641 (g := (s := deque()).appendleft)((0, arepeat(iter_to_agen(it), 1))); f = s.popleft
642 while s:
643 l, n = N = f()
644 if l > levels:
645 async for i in n: yield i
646 continue
647 async for _ in n:
648 if isinstance(_, base_typ): yield _
649 else:
650 try: t = iter_to_agen(_); g((l+1, t)); g(N); break
651 except TypeError: yield _
[docs]
652def afirsttrue(it, default=_NO_DEFAULT, pred=None, await_pred=False): return anext(afilter(pred, it, await_pred), *H.filter_out(default, s=_NO_DEFAULT))
[docs]
653def afirstfalse(it, default=_NO_DEFAULT, pred=None, await_pred=False): return anext(afilterfalse(pred, it, await_pred), *H.filter_out(default, s=_NO_DEFAULT))
[docs]
654async def aprepend(val, it):
655 yield val
656 async for i in iter_to_agen(it): yield i
[docs]
657async def aappend(val, it):
658 async for i in iter_to_agen(it): yield i
659 yield val
[docs]
660async def awrap(it, start, end):
661 yield start
662 async for i in iter_to_agen(it): yield i
663 yield end
[docs]
664async def arandom_product(*a, n=1, _=_randinst.choice):
665 async for i in ancycles(amap(to_tuple, a, await_=True), n): yield _(i)
[docs]
666async def arandom_combination(it, r, _=_sample):
667 (_ := _(range(len(p := await to_tuple(it))), r)).sort()
668 for i in _: yield p[i]
[docs]
669@aawgenf2agenf
670async def arandom_combination_with_replacement(it, r, _=_randrange): return amap((p := await to_tuple(it)).__getitem__, await asorted(arepeat_func(_, r, len(p))))
[docs]
671async def arandom_permutation(it, r=None, _=_sample):
672 p = await to_tuple(it)
673 if r is None: r = len(p)
674 for i in _(p, r): yield i
[docs]
675async def afirst(it, default=_NO_DEFAULT):
676 async for i in iter_to_agen(it): return i
677 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.afirst called on empty iterable without default value')
678 return default
[docs]
679async def alast(it, default=_NO_DEFAULT, _=H.check_methods):
680 try:
681 if _(it, '__getitem__'): return it[-1]
682 return (await to_list(it)).pop() if (f := getattr(it, '__reversed__', None)) is None else next(f())
683 except (IndexError, TypeError, StopIteration, StopAsyncIteration):
684 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.alast called on empty iterable without default value') from None
685 return default
[docs]
686def anth_or_last(it, n, default=_NO_DEFAULT): return alast(aislice(it, n+1), default)
[docs]
687async def awrapf(it, before=None, after=None):
688 if before is not None:
689 r = before()
690 with A.ignore_typeerrs: await r
691 it = iter_to_agen(it)
692 try:
693 async for i in it: yield i
694 finally:
695 if after is not None:
696 r = after()
697 with A.ignore_typeerrs: await r
[docs]
698def abefore_and_after(pred, it): a, b = tee(it); return acompress(atakewhile(pred, a), azip(b)), b
[docs]
699async def anth_combination(it, r, n):
700 if not 0 <= r <= (l := len(p := await to_tuple(it))): raise IndexError(f'asyncutils.iters.anth_combination: {r=} is out of range')
701 c, k = 1, min(r, l-r)
702 for i in range(1, k+1): c = c*(l-k+i)//i
703 if n < 0: n += c
704 if n < 0 or n >= c: raise IndexError(f'asyncutils.iters.anth_combination: {n=} is out of range')
705 while r:
706 c, l, r = c*r//l, l-1, r-1
707 while n >= c: n -= c; c, l = c*(l-r)//l, l-1
708 yield p[~l]
[docs]
709@aawgenf2agenf
710async def asubslices(it): return astarmap(O.getitem, azip(arepeat(s := await to_tuple(it)), astarmap(slice, acombinations(range(len(s)+1), 2))))
[docs]
711async def arepeat_func(f, n=None, /, *a):
712 async def g(i=A.ignore_typeerrs, _=partial(f, *a)): # noqa: B008
713 r = _()
714 with i: r = await r
715 async for _ in aloops(n): await g()
[docs]
716async def apolynomial_from_roots(roots, _=(1,)):
717 async for r in iter_to_agen(roots): _ = aconvolve(_, (1, -r))
718 async for i in iter_to_agen(_): yield i
[docs]
719@aawgenf2agenf
720async def atranspose(mat): return azip(*await to_list(mat), strict=True)
[docs]
721@aawgenf2agenf
722async def aflatten_tensor(tensor, base_typ=(str, bytes), _=H.check_methods):
723 I = iter_to_agen(tensor)
724 while True:
725 try: v = await anext(I)
726 except StopAsyncIteration: break
727 I = aprepend(v, I)
728 if isinstance(v, base_typ) or not (_(v, '__iter__') or _(v, '__aiter__')): break
729 I = aflatten(I)
730 return I
[docs]
731@aawgenf2agenf
732async def apolynomial_derivative(coeff): return amap(O.mul, r := await to_tuple(coeff), range(len(r)-1, 0))
[docs]
733async def apolynomial_eval(coeff, x):
734 if not (n := len(t := await to_tuple(coeff))): return type(x)(0)
735 return await asumprod(t, areversed(await A.collect(apowers(x), n-1)))
[docs]
736@aawgenf2agenf
737async def areshape(mat, shape):
738 if isinstance(shape, int): return batch(aflatten(mat), shape)
739 d = await anext(shape := iter_to_agen(shape)); return aislice(await A.areduce(batch, areversed(shape), aflatten_tensor(mat), await_=False), d)
740async def _factor_pollard(n):
741 if n == 4: return 2 # noqa: PLR2004
742 async for b in arange(1, n):
743 x = y = 2; d = 1
744 while (d := M.gcd((x := (x*x+b)%n)-(y := ((z := (y*y+b)%n)*z+b)%n), n)) == 1: ...
745 if d != n: return d
746 raise ValueError(f'asyncutils.iters.afactor: internal error: {n} is prime')
747@lru_cache
748def _shift_to_odd(n):
749 if not ((1<<(s := ((n-1)^n).bit_length()-1))*(d := n>>s) == n and d&1 and s > -1): raise ValueError(f'asyncutils.iters.aisprime: internal error: {n} is invalid')
750 return s-1, d
751async def _probable_prime(n, base, _=_shift_to_odd):
752 s, d = _(m := n-1)
753 if (x := pow(base, d, n)) in {1, m}: return True
754 async for _ in aloops(s):
755 if (x := x*x%n) == m: return True
756 return False
[docs]
757async def aisprime(n, s=_small_primes, p=_perfect_test, r=_randrange, f=_probable_prime):
758 if n < 210: return n in s # noqa: PLR2004
759 if not (n&1 and n%3 and n%5 and n%7 and n%11 and n%13 and n%17): return False
760 for l, _ in p:
761 if n < l: break
762 else: _ = arepeat_func(r, 64, 2, n-1)
763 return await aall(amap(partial(f, n), _, await_=True))
[docs]
764async def afactor(n, _=_little_primes, F=_factor_pollard):
765 if n < 1: raise ValueError('asyncutils.iters.afactor: no prime factors')
766 if n == 1: return
767 for p in _:
768 while not n%p: yield p; n //= p
769 if n == 1: return
770 e = (t := [n]).extend
771 for n in t:
772 if n < 44521 or await aisprime(n): yield n # noqa: PLR2004
773 else: e((f := await F(n), n//f))
[docs]
784async def arandom_derangement(it, _=_randinst.shuffle):
785 if (l := len(s := await to_tuple(it))) < 2: # noqa: PLR2004
786 if s: raise ValueError('asyncutils.iters.arandom_derangement: no derangements to choose from')
787 return ()
788 i = tuple(p := list(range(l)))
789 while any(map(O.is_, i, p)): _(p); await to_tuple(it)
790 return O.itemgetter(*p)(s)
[docs]
791@aawgenf2agenf
792async def amatmul(*a):
793 M, N = map(iter_to_agen, a); N = aprepend(t := await to_tuple(await anext(N)), N)
794 return batch(astarmap(asumprod, aproduct(M, atranspose(N)), True), len(t))
[docs]
795@aawgenf2agenf
796async def mat_vec_mul(M, V): return amap(asumprod.__get__(await to_tuple(V)), amap(to_tuple, M, await_=True), await_=True)
[docs]
797async def vecs_eq(u, v, cmpeq=H.check, *, strict=True):
798 try: return await aall(amap(cmpeq, u, v, strict=strict))
799 except ValueError: return False
[docs]
800async def afreivalds(A, B, C, k=None, _r=_randrange): n = len(A := await to_tuple(A)); return await aall(await vecs_eq(mat_vec_mul(A, mat_vec_mul(B, r := await to_tuple(arepeat_func(_r, n, 2)))), mat_vec_mul(C, r), int.__eq__) async for _ in aloops(getcontext().AFREIVALDS_DEFAULT_K if k is None else k))
[docs]
801def basic_collect(*_): return to_list(aislice(*_) if len(_) > 1 else _[0])
[docs]
802async def asubstrings(it):
803 for i in (s := await to_tuple(it)): yield i,
804 async for n in arange(2, c := len(s)+1):
805 async for i in arange(c-n): yield s[i:i+n]
[docs]
806def asubstr_indices(seq, reverse=False):
807 r = range(1, x := len(seq)+1)
808 if reverse: r = reversed(r)
809 return ((seq[i:(j := i+L)], i, j) for L in r async for i in arange(x-L))
[docs]
810def iter_task(it, summaryf=aconsume):
811 async def task(f): t = f(); await summaryf(it); return f()-t
812 return (l := H.get_loop_and_set()).create_task(task(l.time))
[docs]
852async def aintersend(i1, i2):
853 audit('asyncutils.iters.aintersend', H.fullname(i1), H.fullname(i2)); t = None, None; f, g = i1.asend, i2.asend
854 while True: yield (t := tuple(await B.gather(f(t[1]), g(t[0]))))
[docs]
855def asendstream(i1, i2): audit('asyncutils.iters.asendstream', H.fullname(i1), H.fullname(i2)); return amap(i1.asend, i2, await_=True)
[docs]
856async def acat(first=None):
857 audit('asyncutils.iters.acat', first)
858 while True: first = yield first
[docs]
859async def aforever():
860 audit('asyncutils.iters.aforever')
861 while True: yield
862async def _guess(I, l, K, d, e, C, c, a, x, _=_extreme, /): # noqa: PLR0913,PLR0917
863 if l is None and (l := O.length_hint(I, -1)) < 0: raise ValueError('asyncutils.iters.aguessmax or asyncutils.iters.aguessmin called with no estlen argument on iterable not implementing length (hint)')
864 if (r := await _(A.take(aside_effect(c, I := iter_to_agen(I), await_=a), M.ceil(l*A.RECIPROCAL_E)), K, x, C, o := object())) is o:
865 if d is _NO_DEFAULT: raise ValueError('empty (async) iterable passed to asyncutils.iters.aguessmax or asyncutils.iters.aguessmin with no default value')
866 return d
867 k, I = K(r), iterate_with_key(I, K, x)
868 if x: k = await k
869 try:
870 if c is None:
871 async for t, i in I:
872 if C(t, k): return i
873 elif a:
874 async for t, i in I:
875 if C(t, k): return i
876 await c(i)
877 else:
878 async for t, i in I:
879 if C(t, k): return i
880 c(i)
881 return r
882 finally:
883 if not (e is None or e.is_set()): (t := (_ := H.get_loop_and_set().create_task)(aconsume(I))).add_done_callback(lambda _: e.set()); _(e.wait()).add_done_callback(t.cancel)
[docs]
884def aguessmax(it, estlen=None, *, key=None, await_key=False, default=_NO_DEFAULT, finish_event=None, reject_cb=None, await_cb=False, _=_guess): return _(it, estlen, key, default, finish_event, O.gt, reject_cb, await_cb, await_key)
[docs]
885def aguessmin(it, estlen=None, *, key=None, await_key=False, default=_NO_DEFAULT, finish_event=None, reject_cb=None, await_cb=False, _=_guess): return _(it, estlen, key, default, finish_event, O.lt, reject_cb, await_cb, await_key)
[docs]
886async def apowers_of_two(*, init=1, init_shift=0, shift=1):
887 init <<= init_shift
888 while True: yield init; init <<= shift
[docs]
889def apowers(base, start=1): return aprepend(start, arepeat(0)) if base == 0 else arepeat(start) if base == 1 else apowers_of_two(init=base, shift=base.bit_length()-1) if base.bit_count() == 1 else aaccumulate(arepeat(base), O.mul, initial=start)
[docs]
890async def areversed(it, /):
891 try:
892 async for i in iter_to_agen(reversed(it)): yield i
893 except TypeError:
894 f = (it := await to_list(it)).pop
895 while it: yield f()
[docs]
896async def arunlength_encode(it, /):
897 async for k, g in agroupby(it): yield k, ailen(g)
[docs]
898def arunlength_decode(it, /): return aflatten(astarmap(arepeat, it))
899async def _dft(a, i=False, /): return await to_tuple(A.take(apowers(M.e**((1 if i else -1)*1j*M.tau/(N := len(a := await to_tuple(a))))), N)), N, a
[docs]
900async def adft(a, /, _=_dft):
901 R, N, a = await _(a)
902 for k in range(N): yield M.sumprod(a, (R[k*i%N] for i in range(N)))
[docs]
903async def aidft(A, /, _=_dft):
904 R, N, A = await _(A, True)
905 for k in range(N): yield M.sumprod(A, (R[k*n%N] for n in range(N)))/N
906async def _aax(i, k, d, f, a): return (await f(aenumerate(amap(k, i, await_=a)), key=_get1, default=d))[0]
[docs]
907def aargmin(it, key=_identity, default=-1, *, await_key=False, _=_aax): return _(it, key, default, amin, await_key)
[docs]
908def aargmax(it, key=_identity, default=-1, *, await_key=False, _=_aax): return _(it, key, default, amax, await_key)
[docs]
909def aargminmax(it, key=_identity, default=(-1, -1), *, await_key=False, _=_aax): return _(it, key, default, aminmax_keyed, await_key)
910(FirstMisMatch := namedtuple('FirstMisMatch', 'i lrem rrem', module=__name__)).collect_left = (Shorter := namedtuple('Shorter', 'i lrem', module=__name__)).collect = lambda self: A.collect(self.lrem) # ty: ignore[unresolved-attribute]
911FirstMisMatch.collect_right = (Longer := namedtuple('Longer', 'i rrem', module=__name__)).collect = lambda self: A.collect(self.rrem) # ty: ignore[unresolved-attribute]
912FirstMisMatch.remainder = lambda self, strict=False: azip(self.lrem, self.rrem, strict=strict) # ty: ignore[unresolved-attribute]
[docs]
913def longest_common_prefix(*i): return amap(_get0, atakewhile(aall_equal, azip(*i)))
[docs]
914async def diff_with(*_, cmpeq=O.eq):
915 i, j = map(iter_to_agen, _)
916 async for n in acount():
917 try: v = await anext(i)
918 except StopAsyncIteration: return Longer(n, j)
919 try: w = await anext(j)
920 except StopAsyncIteration: return Shorter(n, i)
921 if not cmpeq(v, w): return FirstMisMatch(n, aprepend(v, i), aprepend(w, j))
[docs]
922def cloned(it, _=O.methodcaller('copy')): return amap(_, it)
[docs]
923async def fuse(it, end_at=None, *, keep_end=False, yield_after=_NO_DEFAULT):
924 async for i in iter_to_agen(it):
925 if H.check(i, end_at):
926 if keep_end: yield i
927 break
928 yield i
929 if yield_after is _NO_DEFAULT:
930 while True: yield end_at
931 while True: yield yield_after
[docs]
932def map_windows(it, f, n, *, await_=False, star=True): return (astarmap if star else amap)(f, window(it, n), await_=await_)
[docs]
933def advance_by(it, n): return aconsume(aislice(it, n))
[docs]
934def distribute(n, it): return tuple(aislice(j, i, None, n) for i, j in enumerate(tee(it, n)))
[docs]
935def stagger(it, offsets=(-1, 0, 1), **k): return azip_offset(*tee(it, len(offsets)), offsets=offsets, **k)
[docs]
936@aawgenf2agenf
937async def azip_offset(*i, offsets, longest=False, fillvalue=None): s = [A.AChain(arepeat(fillvalue, -n), j) if n < 0 else aislice(j, n, None) if n > 0 else j async for j, n in azip(i, offsets, strict=True)]; return aziplongest(*s, fillvalue=fillvalue) if longest else azip(*s)
[docs]
938async def at_most_one(it, default=None):
939 try: i = await anext(it := iter_to_agen(it))
940 except StopAsyncIteration:
941 if default is A.RAISE: raise ValueError('asyncutils.iters.at_most_one: empty (async) iterable and no default value specified')
942 return default
943 try: j = await anext(it)
944 except StopAsyncIteration: return i
945 raise A.MoreThanOne(aiter(A.AChain((i, j), it)), f'asyncutils.iters.at_most_one: more than one item in (async) iterable; offending item: {j!r}')
[docs]
946async def counts(it, key=None, await_key=False, _=namedtuple('CountItem', 'count key item', module=__name__)):
947 g = (d := {}).get
948 async for k, i in iterate_with_key(it, key, await_key): d[k] = c = g(k, 0)+1; yield _(c, k, i)
[docs]
949async def lstrip(it, pred=None, await_pred=False):
950 async for k, i in iterate_with_key(it := iter_to_agen(it), pred, await_pred):
951 if not k: yield i; break
952 async for i in it: yield i
[docs]
953async def rstrip(it, pred=None, await_pred=False):
954 a, c = (b := []).append, b.clear
955 async for k, i in iterate_with_key(it, pred, await_pred):
956 if k: a(i); continue
957 for j in b: yield j
958 c(); yield i
[docs]
959def strip(it, pred=None, await_pred=False): return lstrip(rstrip(it, pred, await_pred), pred, await_pred)
[docs]
960async def aichunked(it, n):
961 it, n = iter_to_agen(it), n-1
962 async for i in it: a, b = tee(r := aislice(it, n)); yield A.AChain((i,), r, a); await aconsume(b)
[docs]
963async def product_index(p, *i, repeat=1):
964 r = 0
965 async for e, j in azip(p, ancycles(amap(to_tuple, i, await_=True), repeat), strict=True): r = r*len(j)+j.index(e)
966 return r
[docs]
967async def circular_shifts(it, steps=1):
968 r, n = (b := await to_deque(it)).rotate, len(b)
969 if steps == 0: raise ValueError('asyncutils.iters.circular_shifts: steps must be non-zero')
970 async for _ in arepeat(None, n//M.gcd(n, steps := -steps)): yield tuple(b); r(steps)
[docs]
971async def gray_product(*i, repeat=1):
972 for a in (i := await to_tuple(ancycles(amap(to_tuple, i, await_=True), repeat))):
973 if len(a) < 2: raise ValueError('asyncutils.iters.gray_product: each iterable must have at least two items') # noqa: PLR2004
974 b, f, o = [0]*(c := len(i)), list(range(c+1)), [1]*c
975 while True:
976 yield tuple(i[j][b[j]] for j in range(c))
977 j, f[0] = f[0], 0
978 if j == c: return
979 if (x := b[j]+o[j]) == 0 or x == len(i[j])-1: o[j], f[j], f[k] = -o[j], f[k := j+1], k
980 b[j] = x
[docs]
981async def partial_product(*i, repeat=1):
982 i = await to_tuple(amap(iter, ancycles(amap(to_tuple, i, await_=True), repeat)))
983 try: p = [next(j) for j in i]
984 except StopIteration: return
985 yield tuple(p)
986 async for k, a in aenumerate(i):
987 for p[k] in a: yield tuple(p)
[docs]
988async def partitions(it):
989 async for i in apowerset(arange(1, n := len(s := await to_tuple(it)))): yield [s[j:k] async for j, k in azip_offset(i, i, offsets=(1, 0), longest=True, fillvalue=n)]
[docs]
990def scan(f, s, /, *i, await_=False): return amap(partial(f, deque((s,), 1)), *i, await_=await_)
[docs]
991async def sort_together(its, key_list=(0,), key=None, reverse=False, strict=False, *, await_key=False):
992 key_list = await to_tuple(key_list)
993 if key is None:
994 if await_key: raise ValueError('asyncutils.iters.sort_together: cannot await key function if key is None')
995 a = O.itemgetter(*key_list)
996 elif len(key_list) == 1:
997 def a(z, _=key_list[0]): return key(z[_])
998 else:
999 def a(z, _=O.itemgetter(*key_list)): return key(*_(z))
1000 return list(zip(*await asorted(azip(*await to_list(its), strict=strict), key=a, reverse=reverse, await_key=await_key), strict=strict))
[docs]
1001def atabulate_finite(f, *a, await_=True): return amap(f, range(*a), await_=await_)
[docs]
1002async def pad_using(it, f, size=None, *, await_=True):
1003 i = 0
1004 async for j in iter_to_agen(it): yield j; i += 1
1005 async for j in atabulate(f, i, await_=await_) if size is None else atabulate_finite(f, i, size, await_=await_): yield j
[docs]
1006async def mark_ends(it):
1007 a, f = await anext(it := iter_to_agen(it)), True
1008 async for b in it: yield f, False, a; a, f = b, False
1009 yield f, True, a
[docs]
1010def locate(it, pred=bool, window_size=None, await_pred=False):
1011 if window_size is None: return acompress(acount(), amap(pred, it, await_=await_pred))
1012 if window_size < 1: raise ValueError('asyncutils.iters.locate: window_size must be at least 1')
1013 return acompress(acount(), (pred(*w) async for w in window(it, window_size)))
[docs]
1014async def split_when(it, pred, maxsplit=-1, *, await_pred=False):
1015 if maxsplit == 0: yield await to_list(it); return
1016 it = iter_to_agen(it)
1017 try: c = await anext(it)
1018 except StopAsyncIteration: return
1019 b = [c]
1020 async for i in it:
1021 x = pred(c, i)
1022 if await x if await_pred else x:
1023 yield b
1024 if maxsplit == 1: yield await to_list(aprepend(i, it)); return
1025 b = []; maxsplit -= 1
1026 b.append(i); c = i
1027 yield b
[docs]
1028def increasing_runs(it, typ=None, max_runs=-1): return split_when(it, O.ge if typ is None else typ.__ge__, max_runs)
[docs]
1029def decreasing_runs(it, typ=None, max_runs=-1): return split_when(it, O.le if typ is None else typ.__le__, max_runs)
[docs]
1030async def coalesce(it, f, await_=False):
1031 try: c = await anext(it := iter_to_agen(it))
1032 except StopAsyncIteration: return
1033 async for i in it:
1034 j = f(c, i)
1035 if await_: j = await j
1036 if j is A.NO_COALESCE: yield c; c = i
1037 else: c = j
1038 yield c
[docs]
1039async def duplicates(it, key=None, *, await_key=False):
1040 g = (d := {}).get
1041 async for k, i in iterate_with_key(it, key, await_key):
1042 if (r := g(k)) is None: d[k] = True
1043 elif r: yield i; d[k] = False
[docs]
1044def filter_identical(it, s): return afilter(partial(O.is_not, s), it)
[docs]
1045async def afirsttrue_or_last(it, default=_NO_DEFAULT, pred=None, *, await_pred=False):
1046 async for k, i in iterate_with_key(it, pred, await_pred):
1047 if k: return i
1048 default = i
1049 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.afirsttrue_or_last: called on empty iterable without default value')
1050 return default
[docs]
1051async def afirsttrue_or_first(it, default=_NO_DEFAULT, pred=None, *, await_pred=False):
1052 it = iterate_with_key(it, pred, await_pred)
1053 try: k, x = await anext(it)
1054 except StopAsyncIteration:
1055 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.afirsttrue_or_first: called on empty iterable without default value')
1056 return default
1057 if not k:
1058 async for k, i in it:
1059 if k: return i
1060 return x
[docs]
1061async def afirstfalse_or_last(it, default=_NO_DEFAULT, pred=None, *, await_pred=False):
1062 async for k, i in iterate_with_key(it, pred, await_pred):
1063 if not k: return i
1064 default = i
1065 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.afirstfalse_or_last: called on empty iterable without default value')
1066 return default
[docs]
1067async def afirstfalse_or_first(it, default=_NO_DEFAULT, pred=None, *, await_pred=False):
1068 it = iterate_with_key(it, pred, await_pred)
1069 try: k, x = await anext(it)
1070 except StopAsyncIteration:
1071 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.afirstfalse_or_first: called on empty iterable without default value')
1072 return default
1073 if k:
1074 async for k, i in it:
1075 if not k: return i
1076 return x
[docs]
1077async def aiequals(*i, strict=True):
1078 try: return await aall(amap(aall_equal, azip(*i, strict=strict), await_=True))
1079 except ValueError: return False
[docs]
1080async def map_if_else(pred, it, func, func_else=None, *, await_func=False, await_func_else=False, await_pred=False):
1081 if func_else is None and await_func_else: raise TypeError('asyncutils.iters.map_if_else: func_else cannot be None if await_func_else is True')
1082 async for k, i in iterate_with_key(it, pred, await_pred):
1083 if k: yield (await func(i)) if await_func else func(i)
1084 else: yield i if func_else is None else (await func_else(i)) if await_func_else else func_else(i)
[docs]
1085async def is_sorted(it, key=None, reverse=False, strict=False, *, await_key=False):
1086 if key is None:
1087 if await_key: raise TypeError('asyncutils.iters.is_sorted: key cannot be None if await_key is True')
1088 else: it = amap(key, it, await_=await_key)
1089 a, b = tee(it, maxqsize=2); await anext(b, None)
1090 if reverse: b, a = a, b
1091 return await aall(amap(O.lt, a, b)) if strict else not await aany(amap(O.lt, b, a))
1092async def _am(I, K, d, a, c, /):
1093 I = iterate_with_key(I, K, a)
1094 try: k, i = await anext(I)
1095 except StopAsyncIteration:
1096 if d is _NO_DEFAULT: raise ValueError('asyncutils.iters.all_min or asyncutils.iters.all_max called on empty iterable without default value')
1097 return [d]
1098 a = (m := [i]).append
1099 async for x, i in I:
1100 if x == k: a(i)
1101 elif c(x, k): k, m[:] = x, (i,)
1102 return m
[docs]
1103def all_min(it, key=None, default=_NO_DEFAULT, *, await_key=False, _=_am): return _(it, key, default, await_key, O.lt)
[docs]
1104def all_max(it, key=None, default=_NO_DEFAULT, *, await_key=False, _=_am): return _(it, key, default, await_key, O.gt)
[docs]
1105def iterate_with_key(it, key=None, await_key=False):
1106 if key is None:
1107 if await_key: raise TypeError('asyncutils.iters.iterate_with_key: key cannot be None if await_key is True')
1108 def k(i): return i, i
1109 elif await_key:
1110 async def k(i): return await key(i), i
1111 else:
1112 def k(i): return key(i), i
1113 return amap(k, it, await_=await_key)
[docs]
1114def flat_map(*a, **k): return aflatten(amap(*a, **k))
[docs]
1124def arunning_mean(it): return amap(O.truediv, aaccumulate(it), acount(1))
[docs]
1125@aawgenf2agenf
1126async def apowerset_of_sets(it, *, frozen=True): S = tuple(dict.fromkeys(await to_list(amap(frozenset, azip(it))))); return aflatten(astarmap((frozenset if frozen else set).union, acombinations(S, r)) async for r in arange(len(S)+1))
[docs]
1127async def aserialize(it):
1128 l, it = B.Lock(), iter_to_agen(it)
1129 while True:
1130 async with l: x = await anext(it)
1131 yield x # noqa: RUF070
[docs]
1132async def aonline_sorter(it, key=None, reverse=False, *, await_key=False):
1133 audit('asyncutils.iters.aonline_sorter', id(it)); c = C if reverse else __import__('heapq')
1134 if key is None: it = [(x, i) async for i, x in aenumerate(it)]
1135 else: it = [(k, i, x) async for i, (k, x) in aenumerate(iterate_with_key(it, key, await_key))]
1136 if len(it) < 0x20000: c.heapify(it) # noqa: PLR2004
1137 else:
1138 if (e := getattr(aonline_sorter, 'executor', None)) is None: e = H.create_executor(aonline_sorter)
1139 H.get_loop_and_set().run_in_executor(e, c.heapify, it)
1140 a, b, i = partial(c.heappop, it), partial(c.heappush, it), len(it)
1141 if key is None:
1142 while it:
1143 if (j := (yield a()[1])) is not None: b((j, i)); i += 1
1144 elif await_key:
1145 while it:
1146 if (j := (yield a()[2])) is not None: b((await key(j), i, j)); i += 1
1147 else:
1148 while it:
1149 if (j := (yield a()[2])) is not None: b((key(j), i, j)); i += 1
[docs]
1150async def acount_cycle(it, n=None):
1151 if n is None: c = acount(1)
1152 else:
1153 if n == 0: return
1154 if n < 0: raise ValueError('asyncutils.iters.acount_cycle: n must be non-negative')
1155 c = arange(1, n)
1156 a = (s := []).append
1157 async for i in iter_to_agen(it): yield 0, i; a(i)
1158 async for i in azip(arepeat_each(c, len(s := tuple(s))), acycle(s)): yield i
[docs]
1159def arepeat_each(it, n=2): return aflatten(amap(arepeat, it, arepeat(n)))
[docs]
1160async def arepeat_last(it, default=_NO_DEFAULT):
1161 async for default in iter_to_agen(it): yield default # noqa: PLR1704
1162 if default is _NO_DEFAULT: return
1163 if default is A.RAISE: raise A.ItemsExhausted('asyncutils.iters.arepeat_last: (async) iterable exhausted and ``default`` was :const:`~asyncutils.constants.RAISE`')
1164 while True: yield default
[docs]
1165def aadjacent(pred, it, dist=1, *, await_pred=False):
1166 if dist < 0: raise ValueError('asyncutils.iters.aadjacent: dist must be non-negative')
1167 if pred is None:
1168 if await_pred: raise TypeError('asyncutils.iters.aadjacent: pred cannot be None if await_pred is True')
1169 pred = bool
1170 u, v = tee(it, maxqsize=dist+1); return azip(amap(aany, window(A.AChain(arepeat(False, dist), amap(pred, u, await_=await_pred), arepeat(False, dist)), (dist<<1)+1), await_=True), v)
[docs]
1180def group_from(keys, values, strict=True): return agroupby_transform(azip(keys, values, strict=strict), _get0, _get1)
[docs]
1181async def awindowed_complete(it, n):
1182 if n < 0: raise ValueError('asyncutils.iters.awindowed_complete: n must be non-negative')
1183 if n > (l := len(s := await to_tuple(it))): raise ValueError('asyncutils.iters.awindowed_complete: n cannot exceed the length of the iterable')
1184 async for i in arange(l+1): yield s[:(j := i-n)], s[j:i], s[i:]
[docs]
1185async def aall_unique(it, key=None, *, await_key=False):
1186 a, b = (s := set()).add, (l := []).append
1187 if key is None:
1188 if await_key: raise TypeError('asyncutils.iters.aall_unique: key cannot be None if await_key is True')
1189 it = iter_to_agen(it)
1190 else: it = amap(key, it, await_=await_key)
1191 async for i in it:
1192 try:
1193 if i in s: return False
1194 a(i)
1195 except TypeError:
1196 if i in l: return False
1197 b(i)
1198 return True
[docs]
1199async def anth_product(n, *i, repeat=1):
1200 c = M.prod(N := tuple(map(len, i := await to_tuple(amap(to_tuple, reversed(i), await_=True))*repeat)))
1201 if n < 0: n += c
1202 if not 0 <= n < c: raise IndexError(f'asyncutils.iters.anth_product: {n=} is out of range')
1203 a = (r := []).append
1204 for p, x in zip(i, N, strict=True): n, m = divmod(n, x); a(p[m])
1205 return tuple(reversed(r))
[docs]
1206async def anth_permutation(it, r, n):
1207 l = len(it := await to_list(it))
1208 if r is None: r = l
1209 c = M.perm(l, r)
1210 if n < 0: n += c
1211 if not 0 <= n < c: raise IndexError(f'asyncutils.iters.anth_permutation: {n=} is out of range')
1212 a, q = [0]*r, n*M.factorial(l)//c if r < l else n
1213 async for d in arange(1, n+1):
1214 q, i = divmod(q, d)
1215 if 0 <= (x := n-d) < r: a[x] = i
1216 if q == 0: break
1217 return tuple(map(it.pop, a))
[docs]
1218async def anth_combination_with_replacement(it, r, n):
1219 if r < 0: raise ValueError('asyncutils.iters.anth_combination_with_replacement: r must be non-negative')
1220 l = len(it := await to_tuple(it)); c, i, a = M.comb(l+r-1, r) if l else 0 if r else 1, 0, (b := []).append
1221 if n < 0: n += c
1222 if not 0 <= n < c: raise IndexError(f'asyncutils.iters.anth_combination_with_replacement: {n=} is out of range')
1223 while r:
1224 r -= 1
1225 while l >= 0:
1226 if n < (x := M.comb(l+r-1, r)): break
1227 l, i, n = l-1, i+1, n-x
1228 a(it[i]); await A.yield_to_event_loop
1229 return tuple(b)
[docs]
1230async def adifference(it, func=O.sub, *, yield_initial=True, await_func=False):
1231 u, v = tee(it, maxqsize=1); i = await anext(v)
1232 if yield_initial: yield i
1233 async for i in amap(func, u, v, await_=await_func): yield i
[docs]
1234async def aminmax(*I, default=_NO_DEFAULT):
1235 if not I:
1236 if default is _NO_DEFAULT: raise TypeError('asyncutils.iters.aminmax: expected at least 1 argument, got 0')
1237 return default
1238 if len(I) == 1: I = I[0]
1239 I = iter_to_agen(I)
1240 try: l = h = await anext(I)
1241 except StopAsyncIteration:
1242 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.aminmax: got empty iterable and default not passed') from None
1243 return default
1244 async for x, y in aziplongest(I, I, fillvalue=l): l, h = (min(l, y), max(h, x)) if y < x else (min(l, x), max(h, y))
1245 return l, h
[docs]
1246async def aminmax_keyed(*I, key, await_key=False, default=_NO_DEFAULT):
1247 if not I:
1248 if default is _NO_DEFAULT: raise TypeError('asyncutils.iters.aminmax_keyed: expected at least 1 argument, got 0')
1249 return default
1250 if len(I) == 1: I = I[0]
1251 I = iter_to_agen(I)
1252 try: l = h = await anext(I)
1253 except StopAsyncIteration:
1254 if default is _NO_DEFAULT: raise ValueError('asyncutils.iters.aminmax_keyed: got empty iterable and default not passed') from None
1255 return default
1256 L = key(l)
1257 if await_key: L = await L
1258 H = L
1259 async for (X, x), (Y, y) in aziplongest(*(iterate_with_key(I, key, await_key) for _ in repeat(2)), fillvalue=(l, L)):
1260 if Y < X:
1261 if Y < L: l, L = y, Y
1262 if H < X: h, H = x, X
1263 else:
1264 if X < L: l, L = x, X
1265 if H < Y: h, H = y, Y
1266 return l, h
[docs]
1267@aawgenf2agenf
1268async def aouter_product(f, X, Y, /, *a, **k): return batch(astarmap(C.partial(f, C.Placeholder, C.Placeholder, *a, **k), aproduct(X, Y := await to_tuple(Y)), True), len(Y))
1269P.patch_function_signatures((adifference, 'it, func={}, *, yield_initial=True, await_func=False'), (agroupby_transform, 'it, kf={}, vf=None, rf=None, *, await_kf=False, await_vf=False, await_rf=False'), (tee, 'it, n=2, *, maxqsize=None, put_exc=None, loop=None'), (aonline_sorter, 'it, *, key={}, reverse=False, slow=None'), (aside_effect, 'f, it, /, *, size=None, before=None, after=None'), (apolynomial_from_roots, 'roots'), (adistinct_permutations, 'it, r=None'), (abfs, _ := 'start, neighbours, *, include_start=True'), (adfs, _), (aaccumulate, 'it, func={}, *, initial=None'), (aconvolve, 'signal, kernel'), (aislice, 'it, /, *a'), (ainterleave_randomly, 'its'), (hamming_dist, 'i1, i2, /, cmpeq={}'), (aiter_idx, 'it, value, start=0, stop=None'), (amerge_sorted_by, 'its, *, key={}, await_=False, reverse=False'), (amax, _ := '*it, key={}, default=_NO_DEFAULT'), (amin, _), (asample_weighted, _ := 'it, k, *, rrange={0}, rand={0}'), (asample_l, _), (arandom_combination, _ := 'it, r'), (arandom_combination_with_replacement, _), (asorted, 'it, *, key={}, reverse=False'), (aunique_justseen, _ := 'it, key={}'), (aunique_everseen, _), (agroupby, _), (vecs_eq, 'u, v, cmpeq={}, *, strict=True'), (adft, 'xarr, /'), (aidft, 'Xarr, /'), (aconsume, 'it, n=None'), (aall_equal, 'it, key={}, strict=False'), (aprepend, 'val, it'), (arandom_product, '*a, n=1'), (asattolo, 'it, /'), (aargmin, _ := 'it, key={}, default=-1'), (aargmax, _), (afactor, _ := 'n'), (extract, 'it, indices, fut=None, finish=False'), (alast, 'it, default=_NO_DEFAULT'), (aisprime, _), (aguessmax, _ := 'it, estlen, *, key={}, default=_NO_DEFAULT, finish_event=None'), (aguessmin, _), (aflatten, _ := 'it'), (arandom_derangement, _), (afreivalds, 'A, B, C, k=None'), (basic_collect, 'it, n'), (iter_task, 'it, summaryf={}'), (apadnone, 'it'), (aunzip, 'ait, put_batch=None, fillvalue={}'), (aflatten_tensor, 'tensor, base_typ={}'), (arandom_permutation, 'it, r=None'))
1270del P, _tee_helper, _pp, _pf, _traverse, _aunzip_put, _guess, _aax, _extreme, _buffer_consume, _factor_pollard, _shift_to_odd, _probable_prime, _dft, _little_primes, _randrange, _sample, _small_primes, _perfect_test, _rand, _randinst, _identity, _