From 6d7996924f98bc33e6492e0fa9e36c321d8fb969 Mon Sep 17 00:00:00 2001 From: Devaev Maxim Date: Wed, 5 Jun 2019 20:56:46 +0300 Subject: wait short tasks --- kvmd/aiotools.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) (limited to 'kvmd/aiotools.py') diff --git a/kvmd/aiotools.py b/kvmd/aiotools.py index d0c12537..b164e110 100644 --- a/kvmd/aiotools.py +++ b/kvmd/aiotools.py @@ -26,6 +26,7 @@ import functools import typing from typing import Callable +from typing import Coroutine from typing import TypeVar from typing import Any @@ -41,8 +42,25 @@ def atomic(method: _AtomicF) -> _AtomicF: return typing.cast(_AtomicF, wrapper) -def task(method: Callable[..., Any]) -> Callable[..., asyncio.Task]: +def tasked(method: Callable[..., Any]) -> Callable[..., asyncio.Task]: @functools.wraps(method) async def wrapper(*args: object, **kwargs: object) -> asyncio.Task: - return asyncio.create_task(method(*args, **kwargs)) + return create_short_task(method(*args, **kwargs)) return typing.cast(Callable[..., asyncio.Task], wrapper) + + +_ATTR_SHORT_TASK = "_aiotools_short_task" + + +def create_short_task(coro: Coroutine) -> asyncio.Task: + task = asyncio.create_task(coro) + setattr(task, _ATTR_SHORT_TASK, True) + return task + + +async def gather_short_tasks() -> None: + await asyncio.gather(*[ + task + for task in asyncio.Task.all_tasks() + if getattr(task, _ATTR_SHORT_TASK, False) + ]) -- cgit v1.2.3