1from asyncutils.constants import _NO_DEFAULT
2from asyncutils._internal import py312 as C, helpers as H, log as L, patch as P
3from asyncutils._internal.submodules import channels_all as __all__
4from _functools import partial
5from _weakrefset import WeakSet
6import asyncio as I, asyncutils as A
7from collections import defaultdict, deque, namedtuple
8from itertools import count, repeat, starmap
9from sys import addaudithook, audit
[docs]
10@H.subscriptable
11class Observable(A.LoopContextMixin):
12 __slots__ = '__data', '__evt', '__lock', '__q', '__to_remove'
13 @property
14 def idle(self): return self.__evt.is_set()
15 @property
16 def notifying(self): return not self.idle
[docs]
17 async def notify(self, *a, _ret_exc_=False, **k):
18 if not self: return
19 async with self.__lock:
20 if self.notifying:
21 if (q := self.__q) is None: await self.wait_until_idle()
22 else: await q.put((_ret_exc_, a, k)); return
23 self.__evt.clear()
24 try: await self._notify_helper(_ret_exc_, a, k); await self.handle_notifications()
25 finally: self.__evt.set(); await self.handle_unsubscriptions()
[docs]
26 async def notify_sequential(self, *a, _silent_=False, _persistent_=False, **k):
27 for observer in self.__data.copy():
28 try: yield await observer(*a, **k)
29 except Exception:
30 if _silent_:
31 if _persistent_: continue
32 break
33 if not _persistent_: raise
34 L.exception('asyncutils.channels.Observable: error in observer')
[docs]
35 async def wait_for_next(self, timeout=None, strict=False):
36 async def f(*a, **k): F.set_result((a, k)) # noqa: RUF029
37 F, u = self.make_fut(), self.subscribe_nowait(f)
38 try: return await I.wait_for(F, timeout)
39 finally: u(strict)
[docs]
40 async def wait_until_idle(self, timeout=None): await I.wait_for(self.__evt.wait(), timeout)
[docs]
41 async def subscribe(self, observer): await self.wait_until_idle(); return self.subscribe_nowait(observer)
[docs]
42 async def unsubscribe(self, observer, strict=False): await self.wait_until_idle(); self.unsubscribe_nowait(observer, strict)
[docs]
43 async def handle_notifications(self):
44 while (q := self.__q) is not None:
45 try: await self._notify_helper(*q.get_nowait())
46 except I.QueueEmpty: break
[docs]
47 async def handle_unsubscriptions(self):
48 async with self.__lock: self.__data.difference_update(s := self.__to_remove); s.clear()
49 async def _notify_helper(self, r, a, k): await I.gather(*self.make_multiple(obs(*a, **k) for obs in self.__data.copy()), return_exceptions=r)
50 def __init__(self, init_observers=(), maxsize=0): audit('asyncutils.channels.Observable', maxsize); self.__data, self.__lock, self.__to_remove, self.__q, self.__evt = set(init_observers), I.Lock(), set(), None if maxsize is None else C.Queue(maxsize), I.Event()
[docs]
51 def __iter__(self): yield from self.__data
[docs]
52 async def __setup__(self): A.LoopContextMixin.__init__(self)
[docs]
53 async def __cleanup__(self): await I.gather(self.handle_notifications(), self.handle_unsubscriptions())
[docs]
54 def start_accumulation(self): return self.restart_accumulation() or True if self.__q is None else False
[docs]
55 async def restart_accumulation(self, flush=True):
56 if flush: await self.handle_notifications()
57 self.__q = C.Queue()
[docs]
58 def subscribe_nowait(self, observer): self.__data.add(observer); return partial(self.unsubscribe_nowait, observer)
[docs]
59 def unsubscribe_eventually(self, observer, asap=True):
60 if asap and self.__evt.is_set(): self.unsubscribe_nowait(observer)
61 else: self.__to_remove.add(observer)
[docs]
62 def unsubscribe_nowait(self, observer, strict=False): getattr(self.__data, 'remove' if strict else 'discard')(observer)
[docs]
63 def subscribe_sync_func(self, observer): return self.subscribe_nowait(A.to_async(observer))
[docs]
64 def ntimes(self, observer, n=None):
65 if n is None: n = A.getcontext().OBSERVABLE_DEFAULT_NTIMES_N
66 if n <= 0: raise ValueError('asyncutils.channels.Observable.ntimes: n must be positive')
67 async def wrapper(*a, **k):
68 nonlocal n; await observer(*a, **k); n -= 1 # ty: ignore[unsupported-operator]
69 if n == 0: await self.unsubscribe(wrapper)
70 self.subscribe_nowait(wrapper); return partial(self.unsubscribe_nowait, wrapper)
[docs]
71 def filter(self, pred, ret_exc=False):
72 f = partial((_ := type(self)())._notify_helper, ret_exc)
73 async def filtered(*a, **k):
74 if pred(*a, **k): await f(a, k)
75 self.subscribe_nowait(filtered); return _
[docs]
76 def map(self, transform, ret_exc=False):
77 f = partial((_ := Observable())._notify_helper, ret_exc)
78 async def mapped(*a, **k): await f(*transform(*a, **k))
79 self.subscribe_nowait(mapped); return _
[docs]
80 def debounce(self, delay, ret_exc=False):
81 f = partial((_ := type(self)())._notify_helper, ret_exc); t = None
82 async def debounced(*a, **k):
83 nonlocal t
84 if t is not None: await A.safe_cancel(t)
85 async def notifier():
86 with A.ignore_cancellation: await I.sleep(delay); await f(a, k)
87 t = self.make(notifier())
88 self.subscribe_nowait(debounced); return _
[docs]
89 def throttle(self, interval, ret_exc=False):
90 f, t = partial((_ := type(self)())._notify_helper, ret_exc), 0
91 async def throttled(*a, **k):
92 nonlocal t
93 with A.event_loop.from_flags(0) as l:
94 if (c := l.time())-t >= interval: t = c; await f(a, k)
95 self.subscribe_nowait(throttled); return _
[docs]
96 def buffer(self, count, ret_exc=False):
97 f, b, c = (_ := type(self)())._notify_helper, [], max(1, count)
98 async def buffered(*a, **k):
99 b.append((a, k))
100 if len(b) >= c: await I.gather(*starmap(f, H.copy_and_clear(b)), return_exceptions=ret_exc)
101 self.subscribe_nowait(buffered); return _
[docs]
102 def at_change(self, key=lambda *a, **k: (a, frozenset(k.items())), ret_exc=False):
103 f, l = partial((_ := type(self)())._notify_helper, ret_exc), object()
104 async def distinct(*a, **k):
105 nonlocal l
106 if (c := key(*a, **k)) != l: l = c; await f(a, k)
107 self.subscribe_nowait(distinct); return _
[docs]
108 def fork(self, ret_exc=False): self.subscribe_nowait(partial((_ := type(self)()).notify, _ret_exc_=ret_exc)); return _
[docs]
109 def merge(*obs, ret_exc=False):
110 p = partial((_ := type(obs[0])()).notify, _ret_exc_=ret_exc)
111 for o in obs: o._data.add(p)
112 return _
[docs]
113class EventBus(A.LoopContextMixin):
114 __slots__ = '_auditing', '_handler', '_is_shutdown', '_lock', '_middlewares', '_published', '_publishers', '_sem', '_subscribers', '_tracking', 'auditor', 'name'
115 def __init__(self, name=None, *, handler=None, max_concurrent=None, tracking_stats=False):
116 if max_concurrent is None: max_concurrent = A.getcontext().EVENT_BUS_DEFAULT_MAX_CONCURRENT
117 def auditor(*a, f=self.is_auditing, _=self.sync_start_publish):
118 if f(): _(*a)
119 audit('asyncutils.channels.EventBus', name, id(self)); self.auditor, self._subscribers, self._published, self._middlewares, self._publishers, self.name, self._lock, self._auditing, self._handler, self._sem, self._is_shutdown, self._tracking, s[None] = auditor, (s := defaultdict(WeakSet)), defaultdict(int), [], set(), f'{H.fullname(self)} {name or f'#{__class__.__inc_cnt()}'}', I.Lock(), False, handler or (lambda _: None), I.Semaphore(max_concurrent), False, tracking_stats, WeakSet()
[docs]
120 def raise_for_shutdown(self):
121 if self._is_shutdown: raise A.BusShutDown(f'{self.name} is shutting down')
[docs]
122 def get_event_stats(self):
123 if self._tracking: return self._published.copy()
124 raise A.BusStatsError(f'{self.name} is not tracking event stats')
[docs]
125 def subscribers_for(self, event_type): return self._subscribers[event_type].copy()
[docs]
126 def events(self): (s := set(self._subscribers)).discard(None); return s
[docs]
127 def has_subscribers(self, event_type): return bool(self._subscribers[event_type])
[docs]
128 def is_subscribed(self, subscriber, event_type=_NO_DEFAULT): return any(subscriber in i for i in self._subscribers.values()) if event_type is _NO_DEFAULT else subscriber in self._subscribers.get(event_type, ())
129 @property
130 def total_subscribers(self): return sum(map(len, self._subscribers.values()))
131 @property
132 def wildcards(self): return self.subscribers_for(None)
133 @property
134 def wildcard_count(self): return len(self._subscribers[None])
135 @property
136 def active_tasks(self): return self._sem._value
137 @property
138 def stream_queue(self):
139 if (r := getattr(self, '_stream_queue', None)) is None: self._stream_queue = r = C.Queue()
140 return r
141 @stream_queue.setter
142 def stream_queue(self, val, /): self._stream_queue = val
[docs]
143 def is_auditing(self): return self._auditing
144 auditing = property(is_auditing, lambda self, val, /: (self.start_audit if val else self.stop_audit)())
[docs]
145 def start_audit(self):
146 if not (self._auditing or getattr(a := self.auditor, 'added', False)): audit('asyncutils.channels.EventBus.start_audit', id(self)); addaudithook(a); self._auditing = a.added = True # ty: ignore[unresolved-attribute]
[docs]
147 def stop_audit(self): audit('asyncutils.channels.EventBus.stop_audit', id(self)); self._auditing = False
[docs]
148 def add_middleware(self, middleware): r = len(m := self._middlewares); m.append((middleware, None)); return r
[docs]
149 def remove_middleware(self, cookie, *, result=None, strict=False):
150 r, m[cookie] = (m := self._middlewares)[cookie], None
151 if r:
152 if (F := r[1]).done(): return F.result()
153 F.set_result(result)
154 elif strict: raise ValueError(cookie)
155 return result
[docs]
156 def add_temp_middleware(self, middleware, until): self._middlewares.append((middleware, until))
[docs]
157 @(c := A.dualcontextmanager(use_existing_executor=False, create_executor=False, strict=False))
158 def audit_context(self):
159 o = not self._auditing
160 try:
161 if o: self.start_audit()
162 yield
163 finally:
164 if o: self.stop_audit()
[docs]
165 @c
166 def tracking_context(self, stats_receiver=None):
167 o = not self._tracking
168 try:
169 if o: self.start_tracking()
170 yield
171 finally:
172 if o: self.stop_tracking() if stats_receiver is None else stats_receiver.set_result(self.stop_tracking(True))
[docs]
173 def start_tracking(self): self._tracking = True
[docs]
174 def stop_tracking(self, ret_stats=False): self._tracking = False; return H.copy_and_clear(self._published) if ret_stats else self._published.clear()
[docs]
175 def subscribe(self, subscriber, /, event_type=None): self.raise_for_shutdown(); self._subscribers[event_type].add(subscriber); return subscriber
[docs]
176 def unsubscribe(self, subscriber, /, event_type=None):
177 self.raise_for_shutdown()
178 try: self._subscribers[event_type].remove(subscriber); return True
179 except KeyError: return False
[docs]
180 def on(self, event_type): return partial(self.subscribe, event_type=event_type)
[docs]
181 def subscriber_count(self, event_type): return len(self._subscribers[event_type])
182 async def _publish_helper(self, d, s, I, *_, f=I.gather): await f(*((self._safe_callback(i, d, *_) for i in I) if s else (i(d, *_) for i in I)))
[docs]
183 async def publish(self, event_type, data=None, *, wait=True, **k):
184 p, f = self.sync_start_publish(event_type, data, **k)
185 if not wait: return
186 try:
187 await p
188 if f: raise ExceptionGroup(f'errors occurred in publishing middlewares of {self.name}', f) from None
189 L.info('%s: publishing of event %r succeeded', self.name, event_type); L.debug('final data: %r', data)
190 except TimeoutError: raise A.BusTimeout(f'{self.name}: publishing of event {event_type!r} took too long') from None
191 finally: await A.safe_cancel(p)
[docs]
192 def sync_start_publish(self, event_type, data=None, *, safe=None, timeout=None, chaperone=None):
193 self.raise_for_shutdown(); f = []
194 if safe is None: safe = A.getcontext().EVENT_BUS_PUBLISH_DEFAULT_SAFE
195 async def g(C=(lambda e, /, a=f.extend, b=f.append: a(e.exceptions) if isinstance(e, BaseExceptionGroup) else b(e)) if chaperone is None else chaperone, D=data):
196 for t in self._middlewares:
197 if t is None: continue
198 m, F = t
199 if F is not None and F.done(): continue
200 try:
201 if I.iscoroutine(D := m(event_type, D)): D = await D
202 except A.CRITICAL: raise A.Critical
203 except (ExceptionGroup, Exception) as e: C(e) # noqa: BLE001
204 except BaseException as e: raise A.BusPublishingError(self, m) from e # ty: ignore[invalid-argument-type]
205 U = self._subscribers
206 if self._tracking: self._published[event_type] += 1
207 s, w = (U[_].copy() for _ in (event_type, None))
208 await I.gather((f := partial(self._publish_helper, D, safe))(s), f(w, event_type))
209 (P := self._publishers).add(p := self.make(I.wait_for(g(), timeout))); p.add_done_callback(lambda p, d=P.discard: d(p)); return p, f
[docs]
210 async def wait_for_event(self, event_type, *, timeout=None, condition=lambda _: True):
211 async def handler(d):
212 if F.done(): return
213 if I.iscoroutine(c := condition(d)): c = await c
214 if c: F.set_result(d)
215 return self.make(I.wait_for(await self.subscribe_until(F := self.loop.create_future(), handler, event_type), timeout))
[docs]
216 def subscribe_until(self, fut, subscriber, event_type=None, *, till_permanent=None, _=A.ignore_cancellation.combined(TimeoutError)): # noqa: B008
217 if fut.done(): raise RuntimeError('asyncutils.channels.EventBus.subscribe_until: future is already done')
218 async def f():
219 with _: r = await I.wait_for(fut, till_permanent); self.unsubscribe(subscriber, event_type); return r
220 self.subscribe(subscriber, event_type); return self.make(f())
[docs]
221 async def feed_event(self, *d, timeout=None):
222 if (q := self.stream_queue).full(): L.warning('%s: event stream buffer full', self.name)
223 try: await I.wait_for(q.put(d[0] if len(d) == 1 else d), timeout)
224 except C.QueueShutDown: L.info('%s: event stream is closing', self.name, exc_info=True)
225 except TimeoutError:
226 if q.full(): L.warning('%s: event stream data lost', self.name, exc_info=True); q.get_nowait(); q.put_nowait(d)
[docs]
227 async def event_stream(self, event_type=None, *, timeout=_NO_DEFAULT, item_timeout=_NO_DEFAULT, bufsize=None):
228 self.raise_for_shutdown()
229 if not self._auditing: audit('asyncutils.channels.EventBus.event_stream', id(self), event_type)
230 t = await self.subscribe_until(F := self.loop.create_future(), partial(self.feed_event, timeout=A.getcontext().EVENT_BUS_STREAM_DEFAULT_TIMEOUT if timeout is _NO_DEFAULT else timeout), event_type); self.stream_queue = q = C.Queue(A.getcontext().EVENT_BUS_STREAM_DEFAULT_BUFFER_SIZE if bufsize is None else bufsize)
231 if _NO_DEFAULT is item_timeout: item_timeout = A.getcontext().EVENT_BUS_STREAM_DEFAULT_ITEM_TIMEOUT
232 try:
233 while True: yield await I.wait_for(q.get(), item_timeout)
234 except C.QueueShutDown: L.info('%s: event stream has been shut down', self.name, exc_info=True)
235 except TimeoutError: L.exception('%s: event stream is stopping because of timeout in waiting for item', self.name)
236 finally: F.set_result(None); await t
[docs]
237 async def shutdown(self, immediate=False, *, timeout=None, preserve_stats=False):
238 if self._is_shutdown: return
239 self._is_shutdown, f = True, self._sem.acquire; self.stop_audit(); self._middlewares.clear()
240 self.clear()
241 if not preserve_stats: self.clear_stats()
242 try:
243 async with I.timeout(timeout):
244 self.stream_queue.shutdown(immediate)
245 for _ in repeat(None, self.active_tasks): await f()
246 except TimeoutError: L.exception('%s: shutdown timed out, some tasks may be incomplete', self.name)
247 finally:
248 if p := self._publishers: await A.safe_cancel_batch(p)
249 del self._lock, self._handler, self._sem, self._publishers
[docs]
250 async def handle_exception(self, e):
251 if I.iscoroutine(e := self._handler(e)): await e
[docs]
252 def clear(self, event_type=_NO_DEFAULT): return self._subscribers.clear() if event_type is _NO_DEFAULT else self._subscribers.pop(event_type, None)
[docs]
253 def clear_all(self): self.clear(); self.clear_stats()
[docs]
254 def clear_wildcards(self): return self.clear(None)
[docs]
255 def clear_stats(self): self._published.clear()
256 async def _safe_callback(self, c, d, t=None, i=None):
257 try:
258 async with self._sem:
259 if I.iscoroutine(r := c(*H.filter_out(t, s=_NO_DEFAULT), d)): await I.wait_for(r, i)
260 except TimeoutError: L.warning('%s: callback %s timed out', self.name, H.fullname(c), exc_info=True)
261 except A.CRITICAL: raise A.Critical
262 except BaseException as e: await self.handle_exception(e) # noqa: BLE001
[docs]
263 async def __setup__(self): super().__init__()
[docs]
264 def __cleanup__(self): return self.shutdown(immediate=True)
265 __inc_cnt = count(1).__next__
266 P.patch_classmethod_signatures((_ := lambda _, /, f='#%d', c=count(1).__next__: f%c(), '')); P.patch_method_signatures((__init__, 'name=None, *, handler=None, max_concurrent=128, tracking_stats=False'), (subscribe_until, 'fut, subscriber, event_type=None, *, till_permanent=None')); WILDCARD, _inc_cnt = None, classmethod(_); del _, c # noqa: B008
[docs]
267@H.subscriptable
268class Rendezvous:
269 __slots__ = '_getters', '_lock', '_loop', '_putters', '_task'
270 def __init__(self, *, loop=None, lock=None): self._getters, self._putters, self._loop, self._lock = deque(), deque(), H.get_loop_and_set() if loop is None else loop, I.Lock() if lock is None else lock; self._make_task()
271 async def _maintainer(self):
272 f, g = I.sleep.__get__(A.getcontext().RENDEZVOUS_MAINTENANCE_INTERVAL), self.cleanup
273 while True: await f(); g()
[docs]
274 async def put(self, v, /, *, timeout=None):
275 try: await self.raising_put(v, timeout=timeout); return True
276 except (I.CancelledError, TimeoutError): return False
[docs]
277 async def raising_put(self, v, /, *, timeout): await I.wait_for(await I.shield(self._put_helper(v)), timeout)
[docs]
278 async def get(self, default=_NO_DEFAULT, *, timeout=None, _=100):
279 f = (p := self._putters).popleft
280 while p:
281 v, F = f()
282 if not F.done(): F.set_result(None); return v
283 if timeout is None and default is not _NO_DEFAULT: return default
284 self._getters.append(F := self._loop.create_future())
285 try: return await I.wait_for(F, timeout)
286 except TimeoutError:
287 if default is _NO_DEFAULT: raise
288 return default
[docs]
289 def __length_hint__(self): return len(self._getters)+len(self._putters)
[docs]
290 def state_snapshot(self, _=namedtuple('StateSnapshot', 'num_getters num_putters num_ops idle', module='asyncutils.channels')): self.cleanup(); t = len(self._getters), len(self._putters); return _(*t, sum(t), not any(t))
[docs]
291 def cleanup(self): self._getters, self._putters = deque(F for F in self._getters if not F.done()), deque(t for t in self._putters if not t[1].done())
[docs]
292 async def exchange(self, v, /, *, asap=False):
293 g, f = self._getters, True
294 async with self._lock:
295 while g:
296 if not (F := g.popleft()).done(): break
297 else: g.append(F := self._loop.create_future()); f = False
298 if f: F.set_result(v); return await self.get()
299 await (self._put_helper if asap else self.put)(v); g.appendleft(F); return await F
300 async def _put_helper(self, v, /):
301 g = self._getters
302 async with self._lock:
303 while g:
304 if not (F := g.popleft()).done(): F.set_result(v); break
305 else: self._putters.append((v, F := self._loop.create_future()))
306 return F
[docs]
307 async def reset(self, _=partial(A.safe_cancel_batch, disembowel=True)):
308 async with self._lock: await I.gather(A.safe_cancel_batch(self._getters, disembowel=True), A.safe_cancel_batch(F async for _, F in A.adisembowel(self._putters)))
309 await A.safe_cancel(self._task); self._make_task()
310 def _make_task(self): self._task = self._loop.create_task(self._maintainer())
311 P.patch_method_signatures((reset, ''), (state_snapshot, ''), (_maintainer, ''))
312del P