python-iterutils 0.2.10.2__py3-none-any.whl → 0.3.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.
iterutils/misc.py ADDED
@@ -0,0 +1,1014 @@
1
+ #!/usr/bin/env python3
2
+ # encoding: utf-8
3
+
4
+ __all__ = [
5
+ "iterable", "map", "filter", "reduce", "zip",
6
+ "chain", "chain_from_iterable", "chunked", "foreach", "through",
7
+ "flatten", "collect", "async_group_collect", "iter_unique",
8
+ "async_iter_unique", "wrap_iter", "wrap_aiter", "peek_iter", "peek_aiter",
9
+ "acc_step", "cut_iter", "context", "backgroud_loop", "gen_startup",
10
+ "async_gen_startup", "do_iter", "do_aiter", "bfs_iter", "bfs_gen",
11
+ ]
12
+
13
+ from asyncio import create_task, sleep as async_sleep
14
+ from builtins import map as _map, filter as _filter, zip as _zip
15
+ from collections import defaultdict, deque
16
+ from collections.abc import (
17
+ AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Buffer,
18
+ Callable, Collection, Container, Coroutine, Generator, Iterable,
19
+ Iterator, Mapping, MutableMapping, MutableSet, MutableSequence,
20
+ Sequence, ValuesView,
21
+ )
22
+ from contextlib import asynccontextmanager, contextmanager, ExitStack, AsyncExitStack
23
+ from copy import copy
24
+ from functools import update_wrapper
25
+ from itertools import batched, chain as _chain, pairwise
26
+ from inspect import isawaitable, iscoroutinefunction
27
+ from sys import _getframe
28
+ from _thread import start_new_thread
29
+ from time import sleep, time
30
+ from types import FrameType
31
+ from typing import (
32
+ cast, overload, Any, AsyncContextManager, ContextManager, Literal,
33
+ )
34
+
35
+ from asynctools import (
36
+ async_filter, async_map, async_reduce, async_zip, async_batched,
37
+ ensure_async, ensure_aiter, async_chain, async_chain_from_iterable,
38
+ async_collect, async_foreach, async_through,
39
+ )
40
+ from texttools import format_time
41
+ from undefined import undefined
42
+
43
+
44
+ def _coalesce(vals, default=None):
45
+ for val in vals:
46
+ if val is not None:
47
+ return val
48
+ return default
49
+
50
+
51
+ @overload
52
+ def _get_async(back: int = 2, /, *, default: Literal[False] = False) -> bool:
53
+ ...
54
+ @overload
55
+ def _get_async[T](back: int = 2, /, *, default: T) -> bool | T:
56
+ ...
57
+ def _get_async[T](back: int = 2, /, *, default: Literal[False] | T = False) -> bool | T:
58
+ """往上查找,从最近的调用栈的命名空间中获取 `async_` 的值
59
+ """
60
+ def iter_frams(f: None | FrameType = _getframe(back)):
61
+ while f:
62
+ yield f.f_locals.get("async_")
63
+ f = f.f_back
64
+ return _coalesce(iter_frams(), default)
65
+
66
+
67
+ def map(
68
+ function: None | Callable,
69
+ iterable: Iterable | AsyncIterable,
70
+ /,
71
+ *iterables: Iterable | AsyncIterable,
72
+ threaded: bool = False,
73
+ ):
74
+ """
75
+ """
76
+ if (
77
+ threaded or
78
+ iscoroutinefunction(function) or
79
+ isinstance(iterable, AsyncIterable) or
80
+ any(isinstance(i, AsyncIterable) for i in iterables)
81
+ ):
82
+ if function is None:
83
+ if iterables:
84
+ return async_zip(iterable, *iterables, threaded=threaded)
85
+ elif threaded:
86
+ return ensure_aiter(iterable, threaded=threaded)
87
+ else:
88
+ return iterable
89
+ return async_map(function, iterable, *iterables)
90
+ if function is None:
91
+ if iterables:
92
+ return _zip(iterable, *iterables)
93
+ else:
94
+ return iterable
95
+ return _map(function, iterable, *iterables)
96
+
97
+
98
+ def filter(
99
+ function: None | Callable,
100
+ iterable: Iterable | AsyncIterable,
101
+ /,
102
+ threaded: bool = False,
103
+ ):
104
+ """
105
+ """
106
+ if threaded or iscoroutinefunction(function) or isinstance(iterable, AsyncIterable):
107
+ return async_filter(function, iterable, threaded=threaded)
108
+ return _filter(function, iterable)
109
+
110
+
111
+ def reduce(
112
+ function: Callable,
113
+ iterable: Iterable | AsyncIterable,
114
+ initial: Any = undefined,
115
+ /,
116
+ threaded: bool = False,
117
+ ):
118
+ """
119
+ """
120
+ if threaded or iscoroutinefunction(function) or isinstance(iterable, AsyncIterable):
121
+ return async_reduce(function, iterable, initial, threaded=threaded)
122
+ from functools import reduce
123
+ if initial is undefined:
124
+ return reduce(function, iterable)
125
+ return reduce(function, iterable, initial)
126
+
127
+
128
+ def zip(
129
+ iterable: Iterable | AsyncIterable,
130
+ /,
131
+ *iterables: Iterable | AsyncIterable,
132
+ threaded: bool = False,
133
+ ):
134
+ """
135
+ """
136
+ if (not threaded and
137
+ isinstance(iterable, Iterable) and
138
+ all(isinstance(it, Iterable) for it in iterables)
139
+ ):
140
+ return _zip(iterable, *iterables)
141
+ return async_zip(iterable, *iterables, threaded=threaded)
142
+
143
+
144
+ @overload
145
+ def chain[T](
146
+ iterable: Iterable[T],
147
+ /,
148
+ *iterables: Iterable[T],
149
+ threaded: Literal[False] = False,
150
+ ) -> Iterator[T]:
151
+ ...
152
+ @overload
153
+ def chain[T](
154
+ iterable: Iterable[T],
155
+ /,
156
+ *iterables: Iterable[T] | AsyncIterable[T],
157
+ threaded: Literal[True],
158
+ ) -> AsyncIterator[T]:
159
+ ...
160
+ @overload
161
+ def chain[T](
162
+ iterable: AsyncIterable[T],
163
+ /,
164
+ *iterables: Iterable[T] | AsyncIterable[T],
165
+ threaded: Literal[False, True] = False,
166
+ ) -> AsyncIterator[T]:
167
+ ...
168
+ def chain[T](
169
+ iterable: Iterable[T] | AsyncIterable[T],
170
+ /,
171
+ *iterables: Iterable[T] | AsyncIterable[T],
172
+ threaded: Literal[False, True] = False,
173
+ ) -> Iterator[T] | AsyncIterator[T]:
174
+ if (not threaded and
175
+ isinstance(iterable, Iterable) and
176
+ all(isinstance(it, Iterable) for it in iterables)
177
+ ):
178
+ return _chain(iterable, *iterables) # type: ignore
179
+ return async_chain(iterable, *iterables, threaded=threaded)
180
+
181
+
182
+ @overload
183
+ def chain_from_iterable[T](
184
+ iterables: Iterable[Iterable[T]],
185
+ threaded: bool = False,
186
+ *,
187
+ async_: Literal[False] = False,
188
+ ) -> Iterator[T]:
189
+ ...
190
+ @overload
191
+ def chain_from_iterable[T](
192
+ iterables: (
193
+ AsyncIterable[Iterable[T]] |
194
+ AsyncIterable[AsyncIterable[T]] |
195
+ AsyncIterable[Iterable[T] | AsyncIterable[T]]
196
+ ),
197
+ threaded: bool = False,
198
+ *,
199
+ async_: bool = False,
200
+ ) -> AsyncIterator[T]:
201
+ ...
202
+ @overload
203
+ def chain_from_iterable[T](
204
+ iterables: (
205
+ Iterable[Iterable[T]] |
206
+ Iterable[AsyncIterable[T]] |
207
+ Iterable[Iterable[T] | AsyncIterable[T]] |
208
+ AsyncIterable[Iterable[T]] |
209
+ AsyncIterable[AsyncIterable[T]] |
210
+ AsyncIterable[Iterable[T] | AsyncIterable[T]]
211
+ ),
212
+ threaded: bool = False,
213
+ *,
214
+ async_: Literal[True],
215
+ ) -> AsyncIterator[T]:
216
+ ...
217
+ def chain_from_iterable[T](
218
+ iterables: (
219
+ Iterable[Iterable[T]] |
220
+ Iterable[AsyncIterable[T]] |
221
+ Iterable[Iterable[T] | AsyncIterable[T]] |
222
+ AsyncIterable[Iterable[T]] |
223
+ AsyncIterable[AsyncIterable[T]] |
224
+ AsyncIterable[Iterable[T] | AsyncIterable[T]]
225
+ ),
226
+ threaded: bool = False,
227
+ *,
228
+ async_: Literal[False, True] = False,
229
+ ) -> Iterator[T] | AsyncIterator[T]:
230
+ if async_ or threaded:
231
+ return async_chain_from_iterable(iterables, threaded=threaded)
232
+ return _chain.from_iterable(iterables) # type: ignore
233
+
234
+ setattr(chain, "from_iterable", chain_from_iterable)
235
+
236
+
237
+ @overload
238
+ def chunked[T](
239
+ iterable: Iterable[T],
240
+ n: int = 1,
241
+ /,
242
+ *,
243
+ threaded: Literal[False] = False,
244
+ ) -> Iterator[Sequence[T]]:
245
+ ...
246
+ @overload
247
+ def chunked[T](
248
+ iterable: Iterable[T],
249
+ n: int = 1,
250
+ /,
251
+ *,
252
+ threaded: Literal[True],
253
+ ) -> Iterator[Sequence[T]]:
254
+ ...
255
+ @overload
256
+ def chunked[T](
257
+ iterable: AsyncIterable[T],
258
+ n: int = 1,
259
+ /,
260
+ *,
261
+ threaded: Literal[False, True] = False,
262
+ ) -> AsyncIterator[Sequence[T]]:
263
+ ...
264
+ def chunked[T](
265
+ iterable: Iterable[T] | AsyncIterable[T],
266
+ n: int = 1,
267
+ /,
268
+ *,
269
+ threaded: Literal[False, True] = False,
270
+ ) -> Iterator[Sequence[T]] | AsyncIterator[Sequence[T]]:
271
+ """
272
+ """
273
+ if n < 0:
274
+ n = 1
275
+ if isinstance(iterable, Sequence):
276
+ if n == 1:
277
+ return ((e,) for e in iterable)
278
+ return (iterable[i:j] for i, j in pairwise(range(0, len(iterable)+n, n)))
279
+ elif not threaded and isinstance(iterable, Iterable):
280
+ return batched(iterable, n)
281
+ else:
282
+ return async_batched(iterable, n, threaded=threaded)
283
+
284
+
285
+ def foreach(
286
+ function: Callable,
287
+ iterable: Iterable | AsyncIterable,
288
+ /,
289
+ *iterables: Iterable | AsyncIterable,
290
+ default = None,
291
+ threaded: bool = False,
292
+ ):
293
+ if (threaded or
294
+ isinstance(iterable, AsyncIterable) or
295
+ any(isinstance(it, AsyncIterable) for it in iterables)
296
+ ):
297
+ return async_foreach(
298
+ function,
299
+ iterable,
300
+ *iterables,
301
+ default=default,
302
+ threaded=threaded,
303
+ )
304
+ r = default
305
+ if iterables:
306
+ for args in _zip(iterable, *iterables):
307
+ r = function(*args)
308
+ else:
309
+ for arg in iterable:
310
+ r = function(arg)
311
+ return r
312
+
313
+
314
+ def through(
315
+ iterable: Iterable | AsyncIterable,
316
+ /,
317
+ threaded: bool = False,
318
+ ):
319
+ """
320
+ """
321
+ if threaded or isinstance(iterable, AsyncIterable):
322
+ return async_through(iterable, threaded=threaded)
323
+ for _ in iterable:
324
+ pass
325
+
326
+
327
+ @overload
328
+ def flatten(
329
+ iterable: Iterable,
330
+ /,
331
+ exclude_types: type | tuple[type, ...] = (Buffer, str),
332
+ *,
333
+ threaded: Literal[False] = False,
334
+ ) -> Iterator:
335
+ ...
336
+ @overload
337
+ def flatten(
338
+ iterable: Iterable,
339
+ /,
340
+ exclude_types: type | tuple[type, ...] = (Buffer, str),
341
+ *,
342
+ threaded: Literal[True],
343
+ ) -> Iterator:
344
+ ...
345
+ @overload
346
+ def flatten(
347
+ iterable: AsyncIterable,
348
+ /,
349
+ exclude_types: type | tuple[type, ...] = (Buffer, str),
350
+ *,
351
+ threaded: Literal[False, True] = False,
352
+ ) -> AsyncIterator:
353
+ ...
354
+ def flatten(
355
+ iterable: Iterable | AsyncIterable,
356
+ /,
357
+ exclude_types: type | tuple[type, ...] = (Buffer, str),
358
+ threaded: Literal[False, True] = False,
359
+ ) -> Iterator | AsyncIterator:
360
+ """
361
+ """
362
+ if threaded or not isinstance(iterable, Iterable):
363
+ return async_flatten(iterable, exclude_types, threaded=threaded)
364
+ def gen(iterable):
365
+ for e in iterable:
366
+ if isinstance(e, (Iterable, AsyncIterable)) and not isinstance(e, exclude_types):
367
+ yield from gen(e)
368
+ else:
369
+ yield e
370
+ return gen(iterable)
371
+
372
+
373
+ async def async_flatten(
374
+ iterable: Iterable | AsyncIterable,
375
+ /,
376
+ exclude_types: type | tuple[type, ...] = (Buffer, str),
377
+ threaded: bool = False,
378
+ ) -> AsyncIterator:
379
+ """
380
+ """
381
+ async for e in ensure_aiter(iterable, threaded=threaded):
382
+ if isinstance(e, (Iterable, AsyncIterable)) and not isinstance(e, exclude_types):
383
+ async for e in async_flatten(e, exclude_types, threaded=threaded):
384
+ yield e
385
+ else:
386
+ yield e
387
+
388
+
389
+ @overload
390
+ def collect[K, V](
391
+ iterable: Iterable[tuple[K, V]] | Mapping[K, V],
392
+ /,
393
+ rettype: Callable[[Iterable[tuple[K, V]]], MutableMapping[K, V]],
394
+ *,
395
+ threaded: Literal[False] = False,
396
+ ) -> MutableMapping[K, V]:
397
+ ...
398
+ @overload
399
+ def collect[T](
400
+ iterable: Iterable[T],
401
+ /,
402
+ rettype: Callable[[Iterable[T]], Collection[T]] = list,
403
+ *,
404
+ threaded: Literal[False] = False,
405
+ ) -> Collection[T]:
406
+ ...
407
+ @overload
408
+ def collect[K, V](
409
+ iterable: Iterable[tuple[K, V]] | Mapping[K, V],
410
+ /,
411
+ rettype: Callable[[Iterable[tuple[K, V]]], MutableMapping[K, V]],
412
+ *,
413
+ threaded: Literal[True],
414
+ ) -> Coroutine[Any, Any, MutableMapping[K, V]]:
415
+ ...
416
+ @overload
417
+ def collect[T](
418
+ iterable: Iterable[T],
419
+ /,
420
+ rettype: Callable[[Iterable[T]], Collection[T]] = list,
421
+ *,
422
+ threaded: Literal[True],
423
+ ) -> Coroutine[Any, Any, Collection[T]]:
424
+ ...
425
+ @overload
426
+ def collect[K, V](
427
+ iterable: AsyncIterable[tuple[K, V]],
428
+ /,
429
+ rettype: Callable[[Iterable[tuple[K, V]]], MutableMapping[K, V]],
430
+ *,
431
+ threaded: Literal[False, True] = False,
432
+ ) -> Coroutine[Any, Any, MutableMapping[K, V]]:
433
+ ...
434
+ @overload
435
+ def collect[T](
436
+ iterable: AsyncIterable[T],
437
+ /,
438
+ rettype: Callable[[Iterable[T]], Collection[T]] = list,
439
+ *,
440
+ threaded: Literal[False, True] = False,
441
+ ) -> Coroutine[Any, Any, Collection[T]]:
442
+ ...
443
+ def collect(
444
+ iterable: Iterable | AsyncIterable | Mapping,
445
+ /,
446
+ rettype: Callable[[Iterable], Collection] = list,
447
+ *,
448
+ threaded: Literal[False, True] = False,
449
+ ) -> Collection | Coroutine[Any, Any, Collection]:
450
+ """
451
+ """
452
+ if threaded or not isinstance(iterable, Iterable):
453
+ return async_collect(iterable, rettype, threaded=threaded)
454
+ return rettype(iterable)
455
+
456
+
457
+ @overload
458
+ def group_collect[K, V, C: Container](
459
+ iterable: Iterable[tuple[K, V]],
460
+ mapping: None = None,
461
+ factory: None | C | Callable[[], C] = None,
462
+ threaded: bool = False,
463
+ ) -> dict[K, C]:
464
+ ...
465
+ @overload
466
+ def group_collect[K, V, C: Container, M: MutableMapping](
467
+ iterable: Iterable[tuple[K, V]],
468
+ mapping: M,
469
+ factory: None | C | Callable[[], C] = None,
470
+ threaded: bool = False,
471
+ ) -> M:
472
+ ...
473
+ @overload
474
+ def group_collect[K, V, C: Container](
475
+ iterable: AsyncIterable[tuple[K, V]],
476
+ mapping: None = None,
477
+ factory: None | C | Callable[[], C] = None,
478
+ threaded: bool = False,
479
+ ) -> Coroutine[Any, Any, dict[K, C]]:
480
+ ...
481
+ @overload
482
+ def group_collect[K, V, C: Container, M: MutableMapping](
483
+ iterable: AsyncIterable[tuple[K, V]],
484
+ mapping: M,
485
+ factory: None | C | Callable[[], C] = None,
486
+ threaded: bool = False,
487
+ ) -> Coroutine[Any, Any, M]:
488
+ ...
489
+ def group_collect[K, V, C: Container, M: MutableMapping](
490
+ iterable: Iterable[tuple[K, V]] | AsyncIterable[tuple[K, V]],
491
+ mapping: None | M = None,
492
+ factory: None | C | Callable[[], C] = None,
493
+ threaded: bool = False,
494
+ ) -> dict[K, C] | M | Coroutine[Any, Any, dict[K, C]] | Coroutine[Any, Any, M]:
495
+ """
496
+ """
497
+ if threaded or not isinstance(iterable, Iterable):
498
+ return async_group_collect(iterable, mapping, factory, threaded=threaded)
499
+ if factory is None:
500
+ if isinstance(mapping, defaultdict):
501
+ factory = mapping.default_factory
502
+ elif mapping:
503
+ factory = type(next(iter(ValuesView(mapping))))
504
+ else:
505
+ factory = cast(type[C], list)
506
+ elif callable(factory):
507
+ pass
508
+ elif isinstance(factory, Container):
509
+ factory = cast(Callable[[], C], lambda _obj=factory: copy(_obj))
510
+ else:
511
+ raise ValueError("can't determine factory")
512
+ factory = cast(Callable[[], C], factory)
513
+ if isinstance(factory, type):
514
+ factory_type = factory
515
+ else:
516
+ factory_type = type(factory())
517
+ if issubclass(factory_type, MutableSequence):
518
+ add = getattr(factory_type, "append")
519
+ else:
520
+ add = getattr(factory_type, "add")
521
+ if mapping is None:
522
+ mapping = cast(M, {})
523
+ for k, v in iterable:
524
+ try:
525
+ c = mapping[k]
526
+ except LookupError:
527
+ c = mapping[k] = factory()
528
+ add(c, v)
529
+ return mapping
530
+
531
+
532
+ @overload
533
+ async def async_group_collect[K, V, C: Container](
534
+ iterable: Iterable[tuple[K, V]] | AsyncIterable[tuple[K, V]],
535
+ mapping: None = None,
536
+ factory: None | C | Callable[[], C] = None,
537
+ threaded: bool = False,
538
+ ) -> dict[K, C]:
539
+ ...
540
+ @overload
541
+ async def async_group_collect[K, V, C: Container, M: MutableMapping](
542
+ iterable: Iterable[tuple[K, V]] | AsyncIterable[tuple[K, V]],
543
+ mapping: M,
544
+ factory: None | C | Callable[[], C] = None,
545
+ threaded: bool = False,
546
+ ) -> M:
547
+ ...
548
+ async def async_group_collect[K, V, C: Container, M: MutableMapping](
549
+ iterable: Iterable[tuple[K, V]] | AsyncIterable[tuple[K, V]],
550
+ mapping: None | M = None,
551
+ factory: None | C | Callable[[], C] = None,
552
+ threaded: bool = False,
553
+ ) -> dict[K, C] | M:
554
+ """
555
+ """
556
+ iterable = ensure_aiter(iterable, threaded=threaded)
557
+ if factory is None:
558
+ if isinstance(mapping, defaultdict):
559
+ factory = mapping.default_factory
560
+ elif mapping:
561
+ factory = type(next(iter(ValuesView(mapping))))
562
+ else:
563
+ factory = cast(type[C], list)
564
+ elif callable(factory):
565
+ pass
566
+ elif isinstance(factory, Container):
567
+ factory = cast(Callable[[], C], lambda _obj=factory: copy(_obj))
568
+ else:
569
+ raise ValueError("can't determine factory")
570
+ factory = cast(Callable[[], C], factory)
571
+ if isinstance(factory, type):
572
+ factory_type = factory
573
+ else:
574
+ factory_type = type(factory())
575
+ if issubclass(factory_type, MutableSequence):
576
+ add = getattr(factory_type, "append")
577
+ else:
578
+ add = getattr(factory_type, "add")
579
+ if mapping is None:
580
+ mapping = cast(M, {})
581
+ async for k, v in iterable:
582
+ try:
583
+ c = mapping[k]
584
+ except LookupError:
585
+ c = mapping[k] = factory()
586
+ add(c, v)
587
+ return mapping
588
+
589
+
590
+ @overload
591
+ def iter_unique[T](
592
+ iterable: Iterable[T],
593
+ /,
594
+ seen: None | MutableSet = None,
595
+ *,
596
+ threaded: Literal[False] = False,
597
+ ) -> Iterator[T]:
598
+ ...
599
+ @overload
600
+ def iter_unique[T](
601
+ iterable: Iterable[T],
602
+ /,
603
+ seen: None | MutableSet = None,
604
+ *,
605
+ threaded: Literal[True],
606
+ ) -> AsyncIterator[T]:
607
+ ...
608
+ @overload
609
+ def iter_unique[T](
610
+ iterable: AsyncIterable[T],
611
+ /,
612
+ seen: None | MutableSet = None,
613
+ *,
614
+ threaded: Literal[False, True] = False,
615
+ ) -> AsyncIterator[T]:
616
+ ...
617
+ def iter_unique[T](
618
+ iterable: Iterable[T] | AsyncIterable[T],
619
+ /,
620
+ seen: None | MutableSet = None,
621
+ threaded: Literal[False, True] = False,
622
+ ) -> Iterator[T] | AsyncIterator[T]:
623
+ """
624
+ """
625
+ if threaded or not isinstance(iterable, Iterable):
626
+ return async_iter_unique(iterable, seen, threaded=threaded)
627
+ if seen is None:
628
+ seen = set()
629
+ def gen(iterable):
630
+ add = seen.add
631
+ for e in iterable:
632
+ if e not in seen:
633
+ yield e
634
+ add(e)
635
+ return gen(iterable)
636
+
637
+
638
+ async def async_iter_unique[T](
639
+ iterable: Iterable[T] | AsyncIterable[T],
640
+ /,
641
+ seen: None | MutableSet = None,
642
+ threaded: bool = False,
643
+ ) -> AsyncIterator[T]:
644
+ """
645
+ """
646
+ if seen is None:
647
+ seen = set()
648
+ add = seen.add
649
+ async for e in ensure_aiter(iterable, threaded=threaded):
650
+ if e not in seen:
651
+ yield e
652
+ add(e)
653
+
654
+
655
+ @overload
656
+ def wrap_iter[T](
657
+ iterable: Iterable[T],
658
+ /,
659
+ callprev: None | Callable[[T], Any] = None,
660
+ callnext: None | Callable[[T], Any] = None,
661
+ *,
662
+ threaded: Literal[False] = False,
663
+ ) -> Iterator[T]:
664
+ ...
665
+ @overload
666
+ def wrap_iter[T](
667
+ iterable: Iterable[T],
668
+ /,
669
+ callprev: None | Callable[[T], Any] = None,
670
+ callnext: None | Callable[[T], Any] = None,
671
+ *,
672
+ threaded: Literal[True],
673
+ ) -> AsyncIterator[T]:
674
+ ...
675
+ @overload
676
+ def wrap_iter[T](
677
+ iterable: AsyncIterable[T],
678
+ /,
679
+ callprev: None | Callable[[T], Any] = None,
680
+ callnext: None | Callable[[T], Any] = None,
681
+ threaded: Literal[False, True] = False,
682
+ ) -> AsyncIterator[T]:
683
+ ...
684
+ def wrap_iter[T](
685
+ iterable: Iterable[T] | AsyncIterable[T],
686
+ /,
687
+ callprev: None | Callable[[T], Any] = None,
688
+ callnext: None | Callable[[T], Any] = None,
689
+ threaded: bool = False,
690
+ ) -> Iterator[T] | AsyncIterator[T]:
691
+ """
692
+ """
693
+ if threaded or not isinstance(iterable, Iterable):
694
+ return wrap_aiter(
695
+ iterable,
696
+ callprev=callprev,
697
+ callnext=callnext,
698
+ threaded=threaded,
699
+ )
700
+ if not callable(callprev):
701
+ callprev = None
702
+ if not callable(callnext):
703
+ callnext = None
704
+ def gen():
705
+ for e in iterable:
706
+ callprev and callprev(e)
707
+ yield e
708
+ callnext and callnext(e)
709
+ return gen()
710
+
711
+
712
+ async def wrap_aiter[T](
713
+ iterable: Iterable[T] | AsyncIterable[T],
714
+ /,
715
+ callprev: None | Callable[[T], Any] = None,
716
+ callnext: None | Callable[[T], Any] = None,
717
+ threaded: bool = False,
718
+ ) -> AsyncIterator[T]:
719
+ """
720
+ """
721
+ callprev = ensure_async(callprev, threaded=threaded) if callable(callprev) else None
722
+ callnext = ensure_async(callnext, threaded=threaded) if callable(callnext) else None
723
+ async for e in ensure_aiter(iterable, threaded=threaded):
724
+ callprev and await callprev(e)
725
+ yield e
726
+ callnext and await callnext(e)
727
+
728
+
729
+ @overload
730
+ def peek_iter[T](
731
+ iterable: Iterable[T],
732
+ /,
733
+ threaded: Literal[False] = False,
734
+ ) -> None | Iterator[T]:
735
+ ...
736
+ @overload
737
+ def peek_iter[T](
738
+ iterable: Iterable[T],
739
+ /,
740
+ threaded: Literal[True],
741
+ ) -> Coroutine[Any, Any, None | AsyncIterator[T]]:
742
+ ...
743
+ @overload
744
+ def peek_iter[T](
745
+ iterable: AsyncIterable[T],
746
+ /,
747
+ threaded: Literal[False, True] = False,
748
+ ) -> Coroutine[Any, Any, None | AsyncIterator[T]]:
749
+ ...
750
+ def peek_iter[T](
751
+ iterable: Iterable[T] | AsyncIterable[T],
752
+ /,
753
+ threaded: Literal[False, True] = False,
754
+ ) -> None | Iterator[T] | Coroutine[Any, Any, None | AsyncIterator[T]]:
755
+ if threaded or isinstance(iterable, AsyncIterable):
756
+ return peek_aiter(iterable, threaded=threaded)
757
+ try:
758
+ it = iter(iterable)
759
+ first = next(it)
760
+ return chain((first,), it)
761
+ except StopIteration:
762
+ return None
763
+
764
+
765
+ async def peek_aiter[T](
766
+ iterable: Iterable[T] | AsyncIterable[T],
767
+ /,
768
+ threaded: Literal[False, True] = False,
769
+ ) -> None | AsyncIterator[T]:
770
+ try:
771
+ it = ensure_aiter(iterable, threaded=threaded)
772
+ first = await anext(it)
773
+ return async_chain((first,), it)
774
+ except StopAsyncIteration:
775
+ return None
776
+
777
+
778
+ def acc_step(
779
+ start: int,
780
+ stop: None | int = None,
781
+ step: int = 1,
782
+ ) -> Iterator[tuple[int, int, int]]:
783
+ """
784
+ """
785
+ if stop is None:
786
+ start, stop = 0, start
787
+ for i in range(start + step, stop, step):
788
+ yield start, (start := i), step
789
+ if start != stop:
790
+ yield start, stop, stop - start
791
+
792
+
793
+ def cut_iter(
794
+ start: int,
795
+ stop: None | int = None,
796
+ step: int = 1,
797
+ ) -> Iterator[tuple[int, int]]:
798
+ """
799
+ """
800
+ if stop is None:
801
+ start, stop = 0, start
802
+ for start in range(start + step, stop, step):
803
+ yield start, step
804
+ if start != stop:
805
+ yield stop, stop - start
806
+
807
+
808
+ @overload
809
+ def context[T](
810
+ func: Callable[..., T],
811
+ *ctxs: ContextManager,
812
+ async_: Literal[False],
813
+ ) -> T:
814
+ ...
815
+ @overload
816
+ def context[T](
817
+ func: Callable[..., T] | Callable[..., Awaitable[T]],
818
+ *ctxs: ContextManager | AsyncContextManager,
819
+ async_: Literal[True],
820
+ ) -> Coroutine[Any, Any, T]:
821
+ ...
822
+ @overload
823
+ def context[T](
824
+ func: Callable[..., T] | Callable[..., Awaitable[T]],
825
+ *ctxs: ContextManager | AsyncContextManager,
826
+ async_: None = None,
827
+ ) -> T | Coroutine[Any, Any, T]:
828
+ ...
829
+ def context[T](
830
+ func: Callable[..., T] | Callable[..., Awaitable[T]],
831
+ *ctxs: ContextManager | AsyncContextManager,
832
+ async_: None | Literal[False, True] = None,
833
+ ) -> T | Coroutine[Any, Any, T]:
834
+ """
835
+ """
836
+ if async_ is None:
837
+ if iscoroutinefunction(func):
838
+ async_ = True
839
+ else:
840
+ async_ = _get_async()
841
+ if async_:
842
+ async def call():
843
+ args: list = []
844
+ add_arg = args.append
845
+ with ExitStack() as stack:
846
+ async with AsyncExitStack() as async_stack:
847
+ enter = stack.enter_context
848
+ async_enter = async_stack.enter_async_context
849
+ for ctx in ctxs:
850
+ if isinstance(ctx, AsyncContextManager):
851
+ add_arg(await async_enter(ctx))
852
+ else:
853
+ add_arg(enter(ctx))
854
+ ret = func(*args)
855
+ if isawaitable(ret):
856
+ ret = await cast(Awaitable, ret)
857
+ return ret
858
+ return call()
859
+ else:
860
+ with ExitStack() as stack:
861
+ return func(*map(stack.enter_context, ctxs)) # type: ignore
862
+
863
+
864
+ @overload
865
+ def backgroud_loop(
866
+ call: None | Callable = None,
867
+ /,
868
+ interval: int | float = 0.05,
869
+ *,
870
+ async_: Literal[False],
871
+ ) -> ContextManager:
872
+ ...
873
+ @overload
874
+ def backgroud_loop(
875
+ call: None | Callable = None,
876
+ /,
877
+ interval: int | float = 0.05,
878
+ *,
879
+ async_: Literal[True],
880
+ ) -> AsyncContextManager:
881
+ ...
882
+ @overload
883
+ def backgroud_loop(
884
+ call: None | Callable = None,
885
+ /,
886
+ interval: int | float = 0.05,
887
+ *,
888
+ async_: None = None,
889
+ ) -> ContextManager | AsyncContextManager:
890
+ ...
891
+ def backgroud_loop(
892
+ call: None | Callable = None,
893
+ /,
894
+ interval: int | float = 0.05,
895
+ *,
896
+ async_: None | Literal[False, True] = None,
897
+ ) -> ContextManager | AsyncContextManager:
898
+ """
899
+ """
900
+ if async_ is None:
901
+ if iscoroutinefunction(call):
902
+ async_ = True
903
+ else:
904
+ async_ = _get_async()
905
+ use_default_call = not callable(call)
906
+ if use_default_call:
907
+ start = time()
908
+ def call():
909
+ print(f"\r\x1b[K{format_time(time() - start)}", end="")
910
+ def run():
911
+ while running:
912
+ try:
913
+ yield call
914
+ except Exception:
915
+ pass
916
+ if interval > 0:
917
+ if async_:
918
+ yield async_sleep(interval)
919
+ else:
920
+ sleep(interval)
921
+ running = True
922
+ if async_:
923
+ @asynccontextmanager
924
+ async def actx():
925
+ nonlocal running
926
+ try:
927
+ task = create_task(run())
928
+ yield task
929
+ finally:
930
+ running = False
931
+ task.cancel()
932
+ if use_default_call:
933
+ print("\r\x1b[K", end="")
934
+ return actx()
935
+ else:
936
+ @contextmanager
937
+ def ctx():
938
+ nonlocal running
939
+ try:
940
+ yield start_new_thread(run, ())
941
+ finally:
942
+ running = False
943
+ if use_default_call:
944
+ print("\r\x1b[K", end="")
945
+ return ctx()
946
+
947
+
948
+ def gen_startup[**Args, G: Generator](func: Callable[Args, G], /):
949
+ def wrapper(*args: Args.args, **kwds: Args.kwargs) -> G:
950
+ gen = func(*args, **kwds)
951
+ next(gen)
952
+ return gen
953
+ return update_wrapper(wrapper, func)
954
+
955
+
956
+ def async_gen_startup[**Args, G: AsyncGenerator](func: Callable[Args, G], /):
957
+ async def wrapper(*args: Args.args, **kwds: Args.kwargs) -> G:
958
+ gen = func(*args, **kwds)
959
+ await anext(gen)
960
+ return gen
961
+ return update_wrapper(wrapper, func)
962
+
963
+
964
+ def do_iter[T](
965
+ func: Callable[[], T] | Iterable[T],
966
+ /,
967
+ sentinel=undefined,
968
+ sentinel_excs: type[BaseException] | tuple[type[BaseException], ...] = (),
969
+ ) -> Iterator[T]:
970
+ try:
971
+ yield from iter(func if callable(func) else iter(func).__next__, sentinel)
972
+ except sentinel_excs:
973
+ pass
974
+
975
+
976
+ async def do_aiter[T](
977
+ func: Callable[[], T] | Callable[[], Awaitable[T]] | Iterable[T] | AsyncIterable[T],
978
+ /,
979
+ sentinel=undefined,
980
+ sentinel_excs: type[BaseException] | tuple[type[BaseException], ...] = (),
981
+ ) -> AsyncIterator[T]:
982
+ if callable(func):
983
+ func = ensure_async(func)
984
+ else:
985
+ func = ensure_aiter(func).__anext__
986
+ try:
987
+ while True:
988
+ v = await func()
989
+ if v is sentinel:
990
+ break
991
+ yield v
992
+ except StopAsyncIteration:
993
+ pass
994
+ except sentinel_excs:
995
+ pass
996
+
997
+
998
+ def bfs_iter[T](*initials: T) -> tuple[Iterator[T], Callable[[T], None]]:
999
+ dq = deque(initials)
1000
+ return do_iter(dq.popleft, sentinel_excs=IndexError), dq.append
1001
+
1002
+
1003
+ @gen_startup
1004
+ def bfs_gen[T](*initials) -> Generator[None | T, T | None, None]:
1005
+ dq = deque(initials)
1006
+ push, pop = dq.append, dq.popleft
1007
+ try:
1008
+ p = yield None
1009
+ while True:
1010
+ p = yield (pop() if p is None else push(p))
1011
+ except IndexError:
1012
+ pass
1013
+
1014
+ # TODO: 这个模块添加了很多不必要的函数,需要进行移除