async def catch( source: AsyncIterable[Any], exc_class: Type[BaseException], on_exc: Callable[[BaseException], Stream] = None) -> AsyncGenerator[Any, None]: """ Catch an exception and then switch to the next stream `on_exc` or gracefully terminate, when no stream is given. Parameters ---------- source: AsyncIterable exc_class: BaseException type The exception to catch on_exc: Callable A function, that takes an exception and returns a Stream. Yields ------- Any The results from the source stream or the `on_exc` stream. """ try: async with streamcontext(source) as streamer: async for item in streamer: yield item except exc_class as exc: if on_exc is not None: async with on_exc(exc).stream() as streamer: async for item in streamer: yield item else: yield stream.empty()
async def accumulate_by_key(source: Stream, func=op.add, key=lambda x: None, initializer=lambda x: 0): """Generate a series of accumulated sums (or other binary function) from an asynchronous sequence. If ``initializer`` is present, it is placed before the items of the sequence in the calculation, and serves as a default when the sequence is empty. """ accumulated = {} if not callable(initializer): _init = initializer initializer = lambda x: _init iscorofunc = asyncio.iscoroutinefunction(func) iscoroinitializer = asyncio.iscoroutinefunction(initializer) async with streamcontext(source) as streamer: async for item in streamer: id = key(item) if id in accumulated: past = accumulated[id] else: past = initializer(item) if iscoroinitializer: print('initializing') past = await past print('got ' + str(past)) value = func(past, item) if iscorofunc: value = await value accumulated[id] = value yield id, value
async def finally_action( source: AsyncIterable[Any], func: Awaitable[Any] | Callable[[], Any]) -> AsyncGenerator[Any, None]: """ Wrap a try/finally around a stream context. Parameters ---------- source: AsyncIterable func: Awaitable or Callable The function to be called when the entering the finally-block. Yields ------- Any The results from the source stream """ try: async with streamcontext(source) as streamer: async for item in streamer: yield item finally: if inspect.isawaitable(func): await func else: cast(Callable, func)()
async def context( source: AsyncIterable[Any], context_manager: AsyncContextManager, on_enter: Callable[[], Any] = None, on_exit: Callable[[], Any] = None, ) -> AsyncGenerator[Any, None]: """ Iterate a stream within a context. The on_enter and on_exit callbacks can be used to log the status of the stream. Parameters ---------- source: AsyncIterable context_manager: AsyncContextManager The asynchronous context manager that needs to be entered on_enter; Callable A function to be called once the context has been entered on_exit: Callable A function to be called once the context is left. Yields ------- Any The results from the data stream """ async with context_manager: try: if on_enter is not None: on_enter() async with streamcontext(source) as streamer: async for item in streamer: yield item finally: if on_exit is not None: on_exit()
async def retry(source: AsyncIterable[Any], exc_class: Type[BaseException], interval: float = 0) -> AsyncGenerator[Any, None]: """ Retry a datastream if the exception `exc_class` is thrown. Parameters ---------- source: AsyncIterable exc_class: BaseException type The exception class to catch interval: float The time in seconds to wait between retries Yield ------- Any The results from the stream """ timeout: float = 0 loop = asyncio.get_event_loop() while True: try: async with streamcontext(source) as streamer: async for item in streamer: yield item except exc_class as exc: print(f"caught exception of type {type(exc)}") delay = timeout - loop.time() await asyncio.sleep(delay) timeout = loop.time() + interval continue else: return
async def logs(message: types.Message): """ /rec_logs <username> [<from_date>-<to_date>] """ if message.from_user.id not in [ a.user.id for a in await ( await bot.get_chat(settings.chat)).get_administrators() ]: return from_date, to_date = None, datetime.now() try: from_date, to_date = map( dateparser.parse, message.text.split(" ", maxsplit=2)[2].split("-")) except IndexError: pass date_filter = {} if from_date is not None: date_filter["$gte"] = from_date if to_date is not None: date_filter["$lte"] = to_date print(date_filter) try: user = message.text.split(" ") if len(user) <= 1: raise IndexError() key = "from.id" if not user[1].startswith("@") else "from.username" user = user[1] if user == "-": raise IndexError() except IndexError: if message.reply_to_message is None: return user = str(message.reply_to_message.from_user.id) key = "from.id" user = user.lstrip("@") user = int(user) if user.isdigit() else user data = f"Логи для {user}\n\n" async with aiostream.streamcontext( db.find({ key: user, "date": date_filter })) as st: async for item in st: reply = item.get("reply_to_message", {}) data += (f"[{item['date']}, {item['message_id']}] " f"{reply.get('text', '')} -- {item['text']}\n") url = post_to_hastebin(data) await message.reply(url, disable_web_page_preview=True)
async def test_streamcontext(event_loop): with event_loop.assert_cleanup(): xs = stream.range(3) | add_resource.pipe(1) async with streamcontext(xs) as streamer: it = iter(range(3)) async for item in streamer: assert item == next(it) assert event_loop.steps == [1] with event_loop.assert_cleanup(): xs = stream.range(5) | add_resource.pipe(1) async with xs.stream() as streamer: it = iter(range(5)) async for item in streamer: assert item == next(it) assert event_loop.steps == [1]
async def window(source, interval): """Make sure the elements of an asynchronous sequence are separated in time by the given interval. """ def roundend_loop_time(loop): return round_time_to(interval, milliseconds(loop.time())) interval = milliseconds(interval) loop: asyncio.BaseEventLoop = asyncio.get_event_loop() items = [] time = last_time = roundend_loop_time(loop) async with streamcontext(source) as streamer: async for item in streamer: time = roundend_loop_time(loop) if time == last_time: items.append(item) else: yield items + [item] items = [] last_time = time yield items
async def collect(source): result = [] async with streamcontext(source) as streamer: async for item in streamer: result.append(item) return result
async def power(source, exponent): """Raise the elements of an asynchronous sequence to the given power.""" async with streamcontext(source) as streamer: async for item in streamer: yield item ** exponent
async def the_stream(): async with streamcontext(stream(*items)) as streamer: async for item in streamer: await subs.send(item) await subs.close()