1# ruff: noqa: E722
2from asyncutils._internal import compat as Z, helpers as H, log as L, patch as P
3from asyncutils._internal.submodules import base_all as __all__
4from asyncutils.constants import _NO_DEFAULT, RAISE
5from _functools import partial
6import asyncutils as A, asyncio as I
7from enum import IntFlag as E
8from itertools import batched, repeat
9from sys import audit, exc_info
10b, c = H.check_methods, H.fullname
[docs]
11class event_loop: # noqa: N801
12 __reusable, Flags, State = [], E('Flags', ('FLIP_RELEASE_LOOP_ON_FINALIZATION', 'SILENT_ON_FINALIZE', 'NEVER_CLEAR_TASKS_ON_REUSE', 'CLOSE_EXISTING_ON_EXIT', 'SOMETIMES_CONTINUE_ON_EXIT', 'KEEP_CREATED_OPEN_ON_EXIT', 'CANCEL_ALL_TASKS', 'KEEP_LOOP', 'SUPPRESS_RUNTIME_ERRORS', 'FAIL_SILENT', 'DISALLOW_REUSE', 'NO_REUSE', 'NEVER_ENTER', 'ATTEMPT_AENTER', 'SUPPRESS_INNER_EXIT_ON_RUNTIME_ERROR', 'SUPPRESS_INNER_AEXIT_ON_RUNTIME_ERROR'), module=__name__), E('State', ('ENTERED', 'CREATED_LOOP', 'ENTERED_INNER', 'AENTERED_INNER'), module=__name__); __slots__ = '_flags', '_is', '_loop', '_state', '_task'
[docs]
13 def _get_unclosed_loop(self, factory=I.new_event_loop, _=A.IgnoreErrors(AttributeError)): # pragma: no cover
14 if self._flags&(c := self.Flags).NO_REUSE: return factory()
15 p, L = (pool := self.__reusable).pop, None
16 while pool and ((L := p()).is_closed() or L.is_running()): ...
17 if L is None: return factory()
18 if not self._flags&c.NEVER_CLEAR_TASKS_ON_REUSE:
19 with _: L._ready.clear()
20 with _: L._scheduled.clear()
21 return L
[docs]
22 def factory_reset(self): self._flags = self.Flags(A.getcontext().EVENT_LOOP_BASE_FLAGS)
[docs]
23 def clear_flags(self, mask_to_keep=0): self._flags &= mask_to_keep
[docs]
24 def copy_flags(self): return self.from_flags(self._flags)
[docs]
25 def __contains__(self, f, /): return bool(self._flags&(self.Flags[f.upper()] if isinstance(f, str) else f))
[docs]
26 def flags_eq(self, o, /): return self._flags == (o if isinstance(o, int) else o._flags)
[docs]
27 @classmethod
28 def from_flags(cls, flags, /, m=0x10000):
29 if not 0 <= flags < m: raise OverflowError(f'asyncutils.base.event_loop: flags value {flags:#x} has forbidden bits set')
30 r._flags, r._state, r._is = cls.Flags(flags), cls.State(0), f'asyncutils.base.event_loop at {id(r := object.__new__(cls)):#x}'; return r
31 def __new__(cls, /, **k):
32 F, p = A.getcontext().EVENT_LOOP_BASE_FLAGS, k.pop
33 for f, s in cls.Flags.__members__.items():
34 if (x := p(f.lower(), None)) is None: continue
35 if x: F |= s
36 else: F &= ~s
37 if k: A.raise_exc(TypeError, 'asyncutils.base.event_loop: got unexpected keyword arguments; shown below', notes=k) # pragma: no cover
38 return cls.from_flags(F)
[docs]
39 def __hash__(self): return self._flags
[docs]
40 def __enter__(self, _='asyncutils.base.event_loop: context already entered'):
41 q, S = (f := self._flags)&(c := self.Flags).FAIL_SILENT, self.State
42 if (s := self._state)&S.ENTERED: # pragma: no cover
43 if q: return self._loop
44 raise RuntimeError(_)
45 if (l := I._get_running_loop()) is None: I.set_event_loop(l := self._get_unclosed_loop())
46 else: s |= S.CREATED_LOOP
47 if not f&c.NEVER_ENTER and callable(g := getattr(l, '__enter__', None)): # pragma: no cover
48 try: g(); s |= S.ENTERED_INNER
49 except A.CRITICAL: raise A.Critical
50 except BaseException as e:
51 if not q: raise RuntimeError(f'{self._is}: exception occurred while calling __enter__ of associated event loop: {e}') from e # ty: ignore[unresolved-attribute]
52 if f&c.ATTEMPT_AENTER and callable(g := getattr(l, '__aenter__', None)):
53 try: l.run_until_complete(g()); s |= S.AENTERED_INNER
54 except A.CRITICAL: raise A.Critical
55 except BaseException as e:
56 if not q: raise RuntimeError(f'{self._is}: exception occurred while calling __aenter__ of associated event loop: {e}') from e # ty: ignore[unresolved-attribute]
57 self._loop, self._state = l, s+S.ENTERED; return l
[docs]
58 def __exit__(self, t, v, b, /, _m='%s context not entered', _n='%s context not entered with errors passed into __exit__', _i=A.IgnoreErrors(RuntimeError), _l=L): # noqa: C901,PLR0912
59 n, l, a, z, S = self._is, self._loop, not (f := self._flags)&(d := self.Flags).FAIL_SILENT, not f&d.SUPPRESS_RUNTIME_ERRORS, self.State # ty: ignore[unresolved-attribute]
60 if not (s := self._state)&S.ENTERED:
61 if a: raise RuntimeError(_m%n) if v is None else BaseExceptionGroup(_n%n, tuple(A.unnest_reverse(v))).with_traceback(b)
62 return False
63 if f&d.CANCEL_ALL_TASKS: self._task = l.create_task(safe_cancel_batch(I.all_tasks(l)))
64 if not f&d.DISALLOW_REUSE: self.__reusable.append(l)
65 if not ((c := s&S.CREATED_LOOP) and f&d.SOMETIMES_CONTINUE_ON_EXIT):
66 with _i: l.stop()
67 q, r, self._state = t is not None and issubclass(t, RuntimeError), False, s-S.ENTERED
68 if s&S.ENTERED_INNER: # pragma: no cover
69 if callable(g := getattr(l, '__exit__', None)):
70 try: r = g(None, None, None) if f&d.SUPPRESS_INNER_EXIT_ON_RUNTIME_ERROR and q else g(t, v, b)
71 except A.CRITICAL: _l.critical('%s: critical error while calling __exit__ of associated event loop', n, exc_info=True)
72 except RuntimeError:
73 if z: _l.exception('event loop management shenanigans while exiting associated event loop')
74 except:
75 if a: _l.exception('%s: exception occurred while calling __exit__ of associated event loop', n)
76 elif a: _l.error('%s: __enter__ already called but __exit__ is not present', n)
77 if s&S.AENTERED_INNER: # pragma: no cover
78 if callable(g := getattr(l, '__aexit__', None)) and not r:
79 l.stop(); l.run_forever()
80 try: r = l.run_until_complete(g(None, None, None) if f&d.SUPPRESS_INNER_AEXIT_ON_RUNTIME_ERROR and q else g(t, v, b)) or r
81 except A.CRITICAL: _l.critical('%s: critical error while calling __aexit__ of associated event loop', n, exc_info=True)
82 except RuntimeError:
83 if z: _l.exception('runtime error exiting associated event loop')
84 except:
85 if a: _l.exception('%s: exception occurred while calling __aexit__ of associated event loop', n)
86 elif a: _l.error('%s: __aenter__ already called but __aexit__ is not present', n)
87 if f&d.CLOSE_EXISTING_ON_EXIT or not (c or f&d.KEEP_CREATED_OPEN_ON_EXIT):
88 with _i: l.close()
89 I.set_event_loop(None)
90 if not f&d.KEEP_LOOP: del self._loop
91 return r or (q and not z)
[docs]
92 def __del__(self, _f=L.debug, _g=L.warning, _m='%s: garbage-collecting entered context; you are advised to refactor your code', _w='%s: cannot suppress exceptions from within destructor', _d='destroyed %s'): # pragma: no cover
93 b, n = not (f := self._flags)&(c := self.Flags).SILENT_ON_FINALIZE, self._is # ty: ignore[unresolved-attribute]
94 if not self._state&self.State.ENTERED:
95 if b: _f(_d, n)
96 return
97 if b: _g(_m, n)
98 if f&c.FLIP_RELEASE_LOOP_ON_FINALIZATION: self._flags = f^c.DISALLOW_REUSE
99 if self.__exit__(*exc_info()) and b: _g(_w, n)
[docs]
100 def __reduce__(self, /): return self.from_flags, (self._flags,)
101 def __repr__(self, _=c): return f'{_(self)}.from_flags({self._flags:#4x})'
102 P.patch_method_signatures((__enter__, ''), (__exit__, P.exit_sig), (__del__, ''), (_get_unclosed_loop, 'factory={}')); P.patch_classmethod_signatures((from_flags, 'flags, /'), (__new__, f'*, {'={0}, '.join(Flags._member_names_)}={{0}}'))
103def f(n):
104 async def adisembowel(it, /):
105 if callable(p := getattr(it, n, None)):
106 while it: yield p()
107 if callable(p := getattr(it, 'clear', None)) and I.iscoroutine(p := p()): await p
108 else:
109 async for i in iter_to_agen(it): yield i
110 return adisembowel
111adisembowel, adisembowel_left = map(f, ('pop', 'popleft'))
[docs]
112async def safe_cancel_batch(t, /, *, callback=None, disembowel=False, raising=False, _=c):
113 audit('asyncutils.base.safe_cancel_batch', _(t)); a = (l := []).append
114 async for F in (adisembowel if disembowel else iter_to_agen)(t):
115 if not F.done(): F.cancel(); a(F)
116 r = await I.gather(*l, return_exceptions=True)
117 if callback is None: return
118 async def f(a, /, _=callback): return (await x) if I.iscoroutine(x := _(a)) else x
119 L = len(r := await I.gather(*map(f, r), return_exceptions=True))
120 if raising and (E := tuple(A.unnest_reverse(*filter(BaseException.__instancecheck__, r)))): raise BaseExceptionGroup(f'asyncutils.base.safe_cancel_batch: {f'flattened {L} exception (groups)' if len(E) < L else f'collected {L} exceptions'} thrown by callback function {callback!r}', E)
[docs]
121async def iter_to_agen(it, sentinel=_NO_DEFAULT, *, use_existing_executor=None, create_executor=None, strict=None, a=c, b=b, c=H.check, s=H.create_executor, h=H.get_loop_and_set, w=L.debug, d=0x400, _=type('', (), {'__slots__': ('it',), '__init__': lambda self, it: setattr(self, 'it', it), '__bool__': lambda self, _=b: _(self.it, 'send', 'throw', 'close'), '__enter__': lambda self: None, '__exit__': lambda self, t, v, _, /, e=frozenset((StopAsyncIteration, StopIteration)), f=frozenset(('StopIteration interacts badly with generators and cannot be raised into a Future', 'async generator raised StopIteration')): False if t is None else str(v) in f if t is RuntimeError else (((True if (C := getattr(self.it, 'close', None)) is None else C()) if t in e else (True if (T := getattr(self.it, 'throw', None)) is None else T(v))) or True)})): # noqa: ARG005,C901,PLR0912,PLR0913
122 # ruff: disable[ASYNC119]
123 audit('asyncutils.base.iter_to_agen', a(it))
124 if type(it) in Z.s:
125 for i in batched(it, d):
126 for _ in i: yield _
127 await A.yield_to_event_loop
128 return
129 C = A.getcontext()
130 if b(it, '__aiter__') and not (C.ITER_TO_AGEN_DEFAULT_STRICT if strict is None else strict):
131 if sentinel is _NO_DEFAULT:
132 async for _ in it: yield _
133 elif b(it, 'asend', 'athrow', 'aclose'):
134 l = await (_ := it.asend)(None)
135 while not c(l, sentinel): l = await _((yield l))
136 else:
137 async for l in it:
138 if c(l, sentinel): break
139 yield l
140 return
141 elif not b(it, '__iter__'): raise TypeError(f'asyncutils.base.iter_to_agen: cannot iterate over {it!r} synchronously or asynchronously')
142 e, g = None, _(it := iter(it))
143 if create_executor is None: create_executor = C.ITER_TO_AGEN_DEFAULT_MAY_CREATE_EXECUTOR
144 if C.ITER_TO_AGEN_DEFAULT_USE_EXISTING_EXECUTOR if use_existing_executor is None else use_existing_executor:
145 if (e := getattr(iter_to_agen, 'executor', None)) is None:
146 if create_executor: e = s(iter_to_agen)
147 else: w('asyncutils.base.iter_to_agen: no existing executor')
148 elif create_executor: e = s(iter_to_agen, False)
149 with g:
150 if e is None:
151 if g:
152 l = (_ := it.send)(None)
153 while not c(l, sentinel): l = _((yield l)); await A.yield_to_event_loop
154 else:
155 while not c(l := next(it, sentinel), sentinel): yield l; await A.yield_to_event_loop
156 else:
157 def r(*a, _=h().run_in_executor, e=e): return partial(_, e, *a) # noqa: B008
158 if g:
159 l = await (_ := r(it.send))(None)
160 while True:
161 if c(l, sentinel): break
162 l = await _((yield l))
163 else:
164 _ = r(next, it)
165 while True:
166 if c((l := await _()), sentinel): break
167 yield l
168 # ruff: enable[ASYNC119]
[docs]
169def aiter_to_gen(ait, *, use_futures=None, loop=None, strict=None, a=c, b=b, g=H.get_loop_and_set):
170 audit('asyncutils.base.aiter_to_gen', a(ait)); C, e = A.getcontext(), I.futures._chain_future # ty: ignore[unresolved-attribute]
171 if b(ait, '__iter__') and not (C.AITER_TO_GEN_DEFAULT_STRICT if strict is None else strict): return (yield from ait) # noqa: B901
172 if not b(ait, '__aiter__'): raise TypeError(f'asyncutils.base.aiter_to_gen: cannot iterate over {ait!r} synchronously or asynchronously')
173 d = b(ait := aiter(ait), 'asend', 'athrow', 'aclose')
174 with A.ignore_stop_async_iteration:
175 if loop is None: loop = g()
176 if loop.is_running():
177 if not (C.AITER_TO_GEN_DEFAULT_ALLOW_FUTURES if use_futures is None else use_futures): raise RuntimeError(f'asyncutils.base.aiter_to_gen: cannot convert async iterator {ait!r} to sync in running event loop without using futures')
178 def f(*a, f, c=loop.create_task, g=e, t=I.Future): return g(c(f(*a)), F := t()) or F.result()
179 if d:
180 p, x = partial(f, f=ait.asend), None
181 while True: x = yield p(x)
182 else:
183 p = partial(f, f=ait.__anext__)
184 while True: yield p()
185 else:
186 a = loop.run_until_complete
187 if d:
188 p, x = ait.asend, None
189 while True: x = yield a(p(x))
190 else:
191 p = ait.__anext__
192 while True: yield a(p())
[docs]
193async def take(it, n=None, default=_NO_DEFAULT, _=L.debug, m='asyncutils.base.take: ran out of items to take'):
194 if n is None:
195 async for i in iter_to_agen(it): yield i
196 if default is _NO_DEFAULT: return
197 while True: yield default
198 if n == 0: return
199 async for n, i in aenumerate(it, n-1, -1): # noqa: B020,PLR1704
200 yield i
201 if n == 0: return
202 if default is RAISE: raise A.ItemsExhausted(m)
203 if default is _NO_DEFAULT: _(m)
204 else:
205 for _ in repeat(default, n): yield _
[docs]
206async def collect(it, n=None, default=_NO_DEFAULT, _='asyncutils.base.collect: ran out of items to collect'): return [i async for i in take(it, n, default, m=_)]
[docs]
207async def collect_into(out, it, n=None, default=_NO_DEFAULT, _='asyncutils.base.collect_into: ran out of items to collect'):
208 a = out.append
209 async for i in take(it, n, default, m=_): a(i)
[docs]
210async def drop(it, n, *, raising=False, _=L.debug, m='asyncutils.base.drop: ran out of items to drop'):
211 it = iter_to_agen(it)
212 if n == 0:
213 async for i in it: yield i
214 return
215 async for i, j in aenumerate(it):
216 if i == n: yield j; break
217 else:
218 if raising: raise A.ItemsExhausted(m)
219 _(m); return
220 async for j in it: yield j
[docs]
221async def aenumerate(it, start=0, step=1):
222 async for _ in iter_to_agen(it): yield start, _; start += step
223P.patch_function_signatures((safe_cancel_batch, 'batch, /, *, callback=None, disembowel=False, raising=False'), (iter_to_agen, 'it, sentinel={}, *, use_existing_executor=None, create_executor=None, strict=None'), (aiter_to_gen, 'ait, *, use_futures=None, loop=None, strict=None'), (collect, 'it, n=None, default={}'), (collect_into, 'out, it, n=None, default={}'), (take, 'it, n, default={}'), (drop, 'it, n, *, raising=False'))
224yield_to_event_loop, sleep_forever = object.__new__(type('', (), {'__new__': lambda _: yield_to_event_loop, '__await__': (_ := lambda _: (yield)), **dict.fromkeys(('__repr__', '__str__', '__reduce__'), lambda _, r='asyncutils.base.yield_to_event_loop': r)})), I.sleep.__get__(float('inf'))
225(dummy_task := type(_)(_.__code__.replace(co_flags=0x161, co_name=(_ := 'dummy_task'), co_qualname=_), globals())(None)).close()
226del f, _, P, L, b, c, H