# SPDX-License-Identifier: MIT
# Copyright (c) 2019 Intel Corporation
"""
WARNING: concurrent can be much slower for quick tasks. It is best used for long
running concurrent tasks.
"""
import inspect
import asyncio
from collections import UserList
from contextlib import AsyncExitStack
from typing import (
Dict,
Any,
AsyncIterator,
Tuple,
Type,
AsyncContextManager,
Optional,
Set,
)
from .log import LOGGER
[docs]class AsyncContextManagerListContext(UserList):
def __init__(self, parent: "AsyncContextManagerList"):
UserList.__init__(self)
self.parent = parent
self.__stack = None
self.logger = LOGGER.getChild(
"AsyncContextManagerListContext.%s"
% (self.__class__.__qualname__,)
)
async def __aenter__(self):
self.clear()
self.__stack = AsyncExitStack()
await self.__stack.__aenter__()
for item in self.parent.data:
# Equivalent to entering the Object context then calling the object
# to get the ObjectContext and entering that context. We then
# return a list of all the inner contexts
# >>> async with BaseDataFlowObject() as obj:
# >>> async with obj() as ctx:
# >>> clist.append(ctx)
citem = item()
self.logger.debug("Entering context: %r", citem)
self.data.append(await self.__stack.enter_async_context(citem))
return self
async def __aexit__(self, exc_type, exc_value, traceback):
await self.__stack.aclose()
[docs]class AsyncContextManagerList(UserList):
CONTEXT: Type[
AsyncContextManagerListContext
] = AsyncContextManagerListContext
def __init__(self, *args):
UserList.__init__(self, list(args))
self.__stack = None
self.logger = LOGGER.getChild(
"AsyncContextManagerList.%s" % (self.__class__.__qualname__,)
)
def __call__(self) -> "BaseDataFlowFacilitatorObjectContext":
return self.CONTEXT(self)
async def __aenter__(self):
self.__stack = AsyncExitStack()
await self.__stack.__aenter__()
for item in self.data:
self.logger.debug("Entering: %r", item)
await self.__stack.enter_async_context(item)
return self
async def __aexit__(self, exc_type, exc_value, traceback):
await self.__stack.aclose()
@property
def async_exit_stack(self):
return self.__stack
async def concurrently(
work: Dict[asyncio.Task, Any],
*,
errors: str = "strict",
nocancel: Optional[Set[asyncio.Task]] = None,
) -> AsyncIterator[Tuple[Any, Any]]:
# Set up logger
logger = LOGGER.getChild("concurrently")
# Track if first run
first = True
# Set of tasks we are waiting on
tasks = set(work.keys())
# Return when outstanding operations reaches zero
try:
while first or tasks:
first = False
# Wait for incoming events
done, _pending = await asyncio.wait(
tasks, return_when=asyncio.FIRST_COMPLETED
)
for task in done:
# Remove the task from the set of tasks we are waiting for
tasks.remove(task)
# Get the tasks exception if any
exception = task.exception()
if errors == "strict" and exception is not None:
raise exception
if exception is None:
# Remove the compeleted task from work
complete = work[task]
del work[task]
yield complete, task.result()
# Update tasks in case work has been updated by called
tasks = set(work.keys())
else:
logger.debug(
"[%s] Ignoring exception: %s", task, exception
)
finally:
for task in tasks:
if not task.done() and (nocancel is None or task not in nocancel):
task.cancel()
else:
# For tasks which are done but have exceptions which we didn't
# raise, collect their exceptions
task.exception()
[docs]async def aenter_stack(
obj: Any,
context_managers: Dict[str, AsyncContextManager],
call: bool = True,
) -> AsyncExitStack:
"""
Create a :py:class:`contextlib.AsyncExitStack` then go through each key,
value pair in the dict of async context managers. Enter the context of each
async context manager and call setattr on ``obj`` to set the attribute by
the name of ``key`` to the value yielded by the async context manager.
If ``call`` is true then the context entered will be the context returned by
calling each context manager respectively.
"""
stack = AsyncExitStack()
await stack.__aenter__()
if context_managers is not None:
for key, ctxmanager in context_managers.items():
if call:
if inspect.isfunction(ctxmanager):
ctxmanager = ctxmanager.__get__(obj, obj.__class__)
setattr(
obj, key, await stack.enter_async_context(ctxmanager())
)
else:
setattr(obj, key, await stack.enter_async_context(ctxmanager))
return stack
[docs]def context_stacker(
inherit: Type, context_managers: Dict[str, AsyncContextManager]
) -> Type:
"""
Using :func:`aenter_stack`
"""
class ContextStacker(inherit):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self._stack = None
async def __aenter__(self):
await super().__aenter__()
self._stack = await aenter_stack(self, context_managers)
return self
async def __aexit__(self, exc_type, exc_value, traceback):
await super().__aexit__(exc_type, exc_value, traceback)
await self._stack.aclose()
return ContextStacker