1# ruff: noqa: B008,PLR6301 # ty: ignore[invalid-argument-type]
2import asyncutils as A
3from asyncutils.constants import _NO_DEFAULT
4from asyncutils._internal import compat as Z, py312 as D, patch as P
5from asyncutils._internal.helpers import LoopMixinBase, check, get_loop_and_set, fullname
6from asyncutils._internal.log import info
7from asyncutils._internal.submodules import queues_all as __all__
8from _collections import deque
9from _functools import partial
10from abc import ABCMeta, abstractmethod
11from asyncio import Event, QueueEmpty, QueueFull, wait_for
12from itertools import count
13from sys import _getframe, audit
14ignore_qempty, ignore_qfull = map((f := (ignore_qshutdown := A.IgnoreErrors(D.QueueShutDown)).combined), _ := (QueueEmpty, QueueFull))
15ignore_qerrs, f = f(*_), object.__setattr__
16def _wakeup_next(W):
17 P = W.popleft
18 while W:
19 if not (w := P()).done(): w.set_result(None); break
20class Q:
21 exc = A.ForbiddenOperation; __slots__ = '_ms', 'cancel_extend', 'change_get_password', 'change_put_password', 'empty', 'full', 'get', 'get_nowait', 'join', 'put', 'put_nowait', 'qsize', 'shutdown', 'task_done'
22 def __repr__(self): return f'<password-protected queue at {id(self):#x}>'
23 def __new__(cls, /, *a, _=f.__get__, x='pwd_q.'):
24 (f := _(s := super().__new__(cls)))(*next(i := zip(cls.__slots__, a, strict=True)))
25 for n, v in i: v.__qualname__, v.__name__ = x+n, n; f(n, v)
26 return s
27 def __init_subclass__(cls, e=exc('subclass'), /, **_): raise e
28 def _get(self, _=exc('call _get() on')): raise _
29 def _put(self, _=exc('call _put() on')): raise _
30 def _init(self, maxsize, _=exc('call _init() on')): raise _ # noqa: ARG002
31 @property
32 def maxsize(self): return self._ms # ty: ignore[unresolved-attribute]
33 P.patch_method_signatures((_get, ''), (_put, ''), (_init, 'maxsize')); P.patch_classmethod_signatures((__init_subclass__, '**k'), (__new__, 'maxsize, cancel_extend, change_get_password, change_put_password, empty, full, get, get_nowait, join, put, put_nowait, qsize, shutdown, task_done, /'))
[docs]
34def password_queue(password_put=_NO_DEFAULT, password_get=_NO_DEFAULT, maxsize=0, *, protect_get=False, protect_put=True, can_change_get=False, can_change_put=False, priority=False, lifo=False, init_items=(), strict=True, get_from=None, put_from=None, gettyp=object, puttyp=object, _=Q): # noqa: C901,PLR0913,PLR0915
35 audit('asyncutils.queues.password_queue', get_from if protect_get else None, put_from if protect_put else None); C, E, y, z, U, S, m, b = A.getcontext(), A.done_evt(), (G := deque()).append, (P := deque()).append, 0, False, (L := get_loop_and_set()).create_future, object()
36 try: F = _getframe(1)
37 except ValueError: F = None
38 if protect_get:
39 if password_get is _NO_DEFAULT:
40 if F is None or (password_get := F.f_locals.get(get_from := (C.PASSWORD_QUEUE_DEFAULT_GET_FROM if get_from is None else get_from).strip())) is None is (password_get := F.f_globals.get(get_from)): raise A.GetPasswordRetrievalError(get_from)
41 elif get_from is not None: raise TypeError('asyncutils.queues.password_queue: got both get_from and password_get')
42 if not isinstance(password_get, gettyp): raise A.WrongPasswordType(None, password_get, type(password_get), gettyp)
43 if protect_put:
44 if password_put is _NO_DEFAULT:
45 if F is None or (password_put := F.f_locals.get(put_from := (C.PASSWORD_QUEUE_DEFAULT_PUT_FROM if put_from is None else put_from).strip())) is None is (password_put := F.f_globals.get(put_from)): raise A.PutPasswordRetrievalError(put_from)
46 elif put_from is not None: raise TypeError('asyncutils.queues.password_queue: got both put_from and password_put')
47 if not isinstance(password_put, puttyp): raise A.WrongPasswordType(None, password_put, type(password_put), puttyp)
48 def s(p):
49 if not isinstance(p, gettyp): raise A.WrongPasswordType(q, p, type(p), gettyp)
50 if p is not password_get and (strict or not check(p, password_get)): raise A.WrongPassword(q, p)
51 def t(p):
52 if not isinstance(p, puttyp): raise A.WrongPasswordType(q, p, type(p), puttyp)
53 if p is not password_put and (strict or not check(p, password_put)): raise A.WrongPassword(q, p)
54 def u(p):
55 if not protect_get: return
56 if not p: raise A.GetPasswordMissing
57 p, = p; s(p)
58 def v(p):
59 if not protect_put: return
60 if not p: raise A.PutPasswordMissing
61 p, = p; t(p)
62 if priority: g, p = partial((M := Z if lifo else __import__('heapq')).heappop, l := []), partial(M.heappush, l)
63 else: g, p = (l := []).pop if lifo else (l := deque()).popleft, l.append
64 async def get(*p):
65 u(p)
66 while not l:
67 if S: raise D.QueueShutDown
68 F = m()
69 try: y(F); await F
70 except:
71 F.cancel()
72 with A.ignore_valerrs: G.remove(F)
73 if l and not F.cancelled(): _wakeup_next(G)
74 raise
75 return get_nowait(_=b)
76 def get_nowait(*p, _=None):
77 if not l: raise D.QueueShutDown if S else QueueEmpty
78 if _ is not b: u(p)
79 i = g(); _wakeup_next(P); return i
80 async def put(i, /, *p):
81 v(p)
82 while full():
83 if S: raise D.QueueShutDown
84 z(F := m())
85 try: await F
86 except:
87 F.cancel()
88 with A.ignore_valerrs: P.remove(F)
89 if not (full() or F.cancelled()): _wakeup_next(P)
90 raise
91 return put_nowait(i, _=b)
92 def put_nowait(i, /, *P, _=None):
93 if S: raise D.QueueShutDown
94 if full(): raise QueueFull
95 if _ is not b: v(P)
96 p(i); nonlocal U; U += 1; E.clear(); _wakeup_next(G)
97 def change_get_password(opw, npw):
98 if (S and not l) or not can_change_get: return False
99 if not isinstance(npw, gettyp): return False
100 try: s(opw)
101 except A.CRITICAL: raise A.Critical
102 except: return False # noqa: E722
103 nonlocal password_get; password_get = npw; return True
104 def change_put_password(opw, npw):
105 if S or not can_change_put: return False
106 if not isinstance(npw, puttyp): return False
107 try: t(opw)
108 except A.CRITICAL: raise A.Critical
109 except: return False # noqa: E722
110 nonlocal password_put; password_put = npw; return True
111 def task_done():
112 nonlocal U
113 if U == 0: raise ValueError('task_done() called too many times')
114 U -= 1
115 if U == 0: E.set()
116 def shutdown(immediate=False):
117 nonlocal S, U; S = True
118 if immediate:
119 U -= len(l)
120 if U <= 0: U = 0; E.set()
121 l.clear()
122 for d in (G, P):
123 f = d.popleft
124 while d:
125 if not (F := f()).done(): F.set_result(None)
126 q = _(maxsize, lambda msg=None: False, change_get_password, change_put_password, lambda: not l, full := lambda: 0 < maxsize <= len(l), get, get_nowait, A.discard_retval(E.wait), put, put_nowait, lambda: len(l), shutdown, task_done) # noqa: ARG005
127 if init_items:
128 async def extend(f=Z.partial(put, Z.Placeholder, password_put)):
129 async for i in A.iter_to_agen(init_items): await f(i)
130 q.cancel_extend = L.create_task(extend()).cancel # ty: ignore[invalid-assignment]
131 return q
[docs]
132class PotentQueueBase(D.Queue, LoopMixinBase, metaclass=ABCMeta):
[docs]
133 @abstractmethod
134 def _init(self, maxsize): raise NotImplementedError
[docs]
135 @abstractmethod
136 def _get(self): raise NotImplementedError
[docs]
137 @abstractmethod
138 def _put(self, item): raise NotImplementedError
[docs]
139 @abstractmethod
140 def peek_all(self): raise NotImplementedError
[docs]
141 @abstractmethod
142 def qsize(self): raise NotImplementedError
143 def __init__(self, maxsize=0): super().__init__(maxsize); self._event = Event()
[docs]
144 def reset(self): super().__init__(self.maxsize); self._event.clear()
[docs]
145 async def smart_put(self, item, *, timeout=None, raising=True):
146 try: self.put_nowait(item); return True
147 except QueueFull: ...
148 try: await wait_for(self.put(item), timeout)
149 except TimeoutError:
150 if raising: raise
151 return False
[docs]
152 async def smart_get(self, *, timeout=None, default=_NO_DEFAULT):
153 f = default is _NO_DEFAULT
154 try: return self.get_nowait()
155 except D.QueueShutDown:
156 if f: raise
157 return default
158 except QueueEmpty: ...
159 try: return await wait_for(self.get(), timeout)
160 except TimeoutError as e:
161 if f: raise e from None
162 return default
[docs]
163 async def extend(self, it):
164 info(f'extending {fullname(self)} with iterable {it!r}'); f = self.smart_put
165 async for i in A.iter_to_agen(it): await f(i)
[docs]
166 def push(self, item):
167 try:
168 if self.full(): audit(f'{fullname(self)}.push', id(self), item, self.get_nowait())
169 self.put_nowait(item); return True
170 except D.QueueShutDown: return False
[docs]
171 async def drain_persistent(self, max_items=None, timeout=None, _=ignore_qshutdown.combined(TimeoutError)):
172 m, c = abs(max_items or float('inf')), 0; info(f'persistent draining of {fullname(self)} started')
173 with _:
174 while c < m: yield await wait_for(self.get(), timeout); self.task_done(); c += 1 # noqa: ASYNC119
[docs]
175 def drain_until_empty(self, max_items=None):
176 max_items, c, g = abs(max_items or float('inf')), 0, self.get_nowait; info(f'draining of {fullname(self)} started')
177 with ignore_qempty:
178 while c < max_items: yield g(); c += 1
[docs]
179 def drain_into_list(self, max_items=None): return list(self.drain_until_empty(max_items))
[docs]
180 def __iter__(self): return self.drain_until_empty()
[docs]
181 def __aiter__(self): return self.drain_persistent()
[docs]
182 def shutdown(self, immediate=False): self._event.set(); super().shutdown(immediate)
183 def __repr__(self): return f'{fullname(self)}({self.maxsize})'
184 @property
185 def is_shutdown(self): return self._event.is_set()
186 @is_shutdown.setter
187 def is_shutdown(self, val, /): self.shutdown() if val else self.__init__(self.maxsize)
188 @property
189 def can_put_now(self): return not (self.is_shutdown or self.full())
190 @property
191 def can_get_now(self): return not (self.is_shutdown or self.empty())
192 @property
193 def fully_functional(self): return not (self.is_shutdown or self.full() or self.empty())
194 @property
195 def capacity(self): return m if (m := self.maxsize) > 0 else float('inf')
196 @property
197 def remaining_capacity(self): return self.capacity-self.qsize()
198 @property
199 def utilization_rate(self): return self.qsize()/self.maxsize
[docs]
200 def pushpop_nowait(self, item, raising=True):
201 if self.is_shutdown: raise D.QueueShutDown
202 if self.empty():
203 if raising: raise QueueEmpty(f'{fullname(self)}.pushpop_nowait on {item!r} expected non-empty queue with raising=True')
204 return self.put_nowait(item)
205 if self.full():
206 if raising: raise QueueFull(f'{fullname(self)}.pushpop_nowait on {item!r} expected non-full queue with raising=True')
207 r = self.get_nowait(); self.put_nowait(item); return r
208 self.put_nowait(item); return self.get_nowait()
[docs]
209 def poppush_nowait(self, item, raising=True):
210 if self.is_shutdown: raise D.QueueShutDown
211 if self.empty():
212 if raising: raise QueueEmpty(f'{fullname(self)}.pushpop_nowait on {item!r} expected non-empty queue with raising=True')
213 return self.put_nowait(item)
214 r = self.get_nowait(); self.put_nowait(item); return r
[docs]
215 async def pushpop(self, item): await self.put(item); return await self.get()
[docs]
216 async def poppush(self, item): r = await self.get(); await self.put(item); return r
[docs]
217 def clear(self):
218 with ignore_qempty:
219 while True: self.get_nowait()
[docs]
220 @A.dualcontextmanager
221 def transaction(self, _=A.IgnoreErrors(TimeoutError)):
222 audit((s := f'{fullname(self)}.transaction/%s')%'start', i := id(self)); q = self.peek_all()
223 try: yield self
224 except:
225 self.clear(); f = self.put_nowait
226 for _ in q: f(_)
227 raise
228 finally: audit(s%'end', i)
[docs]
229 def empty(self): return self.qsize() == 0
[docs]
230 def __bool__(self): return self.qsize() >= 0
[docs]
231 def map(self, f, stop_when=None, *, lifo=False):
232 audit(f'{fullname(self)}.map', id(self), fullname(f))
233 if stop_when is None:
234 stop_when, E = A.AsyncCallbacksFuture(loop=self.loop), (D.QueueShutDown, QueueEmpty)
235 async def get(g=self.drain_until_empty, /): # noqa: RUF029
236 try:
237 for i in g(): yield i
238 finally: stop_when.set_result(None)
239 else:
240 E = (D.QueueShutDown,)
241 async def get(g=self.get, /):
242 with ignore_qshutdown:
243 while True: yield await g() # noqa: ASYNC119
244 async def feed(q, s, g, /, f=f):
245 try:
246 while True: await q.put(await f(await anext(g)))
247 except E: await A.safe_cancel(s)
248 (s := A.AsyncCallbacksFuture(loop=self.loop)).add_noargs_async_callback(partial(A.safe_cancel, self.make(feed(q := (SmartLifoQueue if lifo else SmartQueue)(self.maxsize), s, get()))))
249 if stop_when: stop_when.add_done_callback(s.cancel)
250 return q
[docs]
251 def starmap(self, f, stop_when=None, *, lifo=False):
252 audit(f'{fullname(self)}.starmap', id(self), fullname(f))
253 if stop_when is None:
254 stop_when, E = A.AsyncCallbacksFuture(loop=self.loop), (D.QueueShutDown, QueueEmpty)
255 async def get(g=self.drain_until_empty, /): # noqa: RUF029
256 try:
257 for i in g(): yield i
258 finally: stop_when.set_result(None)
259 else:
260 E = D.QueueShutDown,
261 async def get(g=self.get, /):
262 while True: yield await g()
263 async def feed(q, s, g, /, f=f):
264 try:
265 async for _ in g: await q.put(await f(*_))
266 except E: await A.safe_cancel(s)
267 (s := A.AsyncCallbacksFuture(loop=self.loop)).add_noargs_async_callback(partial(A.safe_cancel, self.make(feed(q := (SmartLifoQueue if lifo else SmartQueue)(self.maxsize), s, get()))))
268 if stop_when: stop_when.add_done_callback(s.cancel)
269 return q
[docs]
270 def filter(self, pred=bool, *, lifo=False):
271 audit(f'{fullname(self)}.filter', id(self), fullname(pred))
272 q = (SmartLifoQueue if lifo else SmartQueue)(self.maxsize)
273 async def feed(f=self.smart_put, g=q.smart_put, h=self.get, _=pred):
274 with ignore_qshutdown:
275 while True: await (f if _(i := await h()) else g)(i)
276 self.make(feed()); return q
[docs]
277 def enumerate(self, *, lifo=False):
278 audit(f'{fullname(self)}.enumerate', id(self))
279 q = (SmartLifoQueue if lifo else SmartQueue)(self.maxsize)
280 async def feed():
281 i = 0
282 with ignore_qempty:
283 while True: await q.smart_put((i, await self.get())); i += 1
284 self.make(feed()); return q
[docs]
285 def filter_nowait(self, pred=bool, /):
286 f, g, a = (k := []).append, (r := []).append, self.get_nowait
287 with ignore_qempty:
288 while True: (f if pred(i := a()) else g)(i)
289 h, j = self.put_nowait, len(r)
290 for i in k:
291 try: h(i)
292 except QueueFull: g(i)
293 return r, j
[docs]
294 def enumerate_nowait(self, start=0, *, step=1):
295 with ignore_qempty:
296 while True: yield start, self.get_nowait(); start += step
297 P.patch_method_signatures((filter_nowait, 'pred=bool'), (transaction, ''), (drain_persistent, 'max_items=None, timeout=None'))
[docs]
298class SmartQueue(PotentQueueBase):
[docs]
299 def _init(self, maxsize): self.__queue = deque(maxlen=maxsize if maxsize > 0 else None)
[docs]
300 def _get(self): return self.__queue.popleft()
[docs]
301 def _put(self, item): self.__queue.append(item)
[docs]
302 def peek(self):
303 if q := self.__queue: return q[0]
304 raise QueueEmpty
[docs]
305 def peek_all(self): return list(self.__queue)
[docs]
306 def qsize(self): return len(self.__queue)
[docs]
307 def rotate(self, n=1, /): self.__queue.rotate(n)
308 def __bool__(self): return bool(self.__queue)
309 def empty(self): return not self
[docs]
310class SmartLifoQueue(PotentQueueBase):
[docs]
311 def _init(self, maxsize): self.__queue = [] # noqa: ARG002
[docs]
312 def _get(self): return self.__queue.pop()
[docs]
313 def _put(self, item): self.__queue.append(item)
[docs]
314 def peek(self, i=-1, /):
315 s = self.qsize()
316 if i < 0: i += s
317 if 0 <= i < s: return self.__queue[i]
318 raise IndexError(f'asyncutils.queues.SmartLifoQueue: failed to peek item at index {i}')
[docs]
319 def peek_all(self): return self.__queue.copy()
[docs]
320 def qsize(self): return len(self.__queue)
321 def __bool__(self): return bool(self.__queue)
322 def empty(self): return not self
323 def pushpop(self, item): raise NotImplementedError
324 def pushpop_nowait(self, item, raising=True): raise NotImplementedError
[docs]
325class SmartPriorityQueue(PotentQueueBase):
326 def __init__(self, maxsize=0, *, init_items=()): super().__init__(maxsize); self.make(self.start(maxsize, init_items))
[docs]
327 async def start(self, maxsize, init_items): q = await A.collect(I := A.iter_to_agen(init_items), maxsize); import heapq as H; H.heapify(q); self.__get, self.__put, self._unfinished_tasks, self.__queue = partial(H.heappop, q), partial(H.heappush, q), len(q), q; self._finished.clear(); await self.extend(I) # ty: ignore[unresolved-attribute]
[docs]
328 def _init(self, maxsize): ...
[docs]
329 def _get(self): return self.__get()
[docs]
330 def _put(self, item): self.__put(item)
[docs]
331 def peek(self): return self.__queue[0]
[docs]
332 def peek_all(self): return self.__queue.copy()
[docs]
333 def qsize(self): return len(self.__queue)
334 def __bool__(self): return bool(self.__queue)
335 def empty(self): return not self
[docs]
336class UserPriorityQueue(SmartPriorityQueue):
[docs]
337 @classmethod
338 def from_iter_of_tuples(cls, items, maxsize=0, _=SmartPriorityQueue): _.__init__(Q := object.__new__(cls), maxsize, init_items=items); Q.__tiebreak = count(); return Q
339 def __init__(self, maxsize=0, *, init_items=(), init_priority=0): self.__tiebreak = count(); super().__init__(maxsize, init_items=((init_priority, self._tiebreak, j) async for j in A.iter_to_agen(init_items)))
340 @property
341 def _tiebreak(self): return next(self.__tiebreak)
[docs]
342 def put_nowait(self, item, priority=0): super().put_nowait((priority, self._tiebreak, item))
[docs]
343 def put(self, item, priority=0): return super().put((priority, self._tiebreak, item))
[docs]
344 def get_nowait(self): return super().get_nowait()[-1]
[docs]
345 async def get(self): return (await super().get())[-1]
346del f, _, Q