Skip to content

Instantly share code, notes, and snippets.

@abersheeran
Last active January 22, 2021 04:26
Show Gist options
  • Save abersheeran/bd2372bb35fec859d7fca453ca5f7826 to your computer and use it in GitHub Desktop.
Save abersheeran/bd2372bb35fec859d7fca453ca5f7826 to your computer and use it in GitHub Desktop.
import asyncio
from asyncio import AbstractEventLoop, Task, Future
from typing import Set, Optional, Coroutine
def onlyfirst(*coros: Coroutine, loop: Optional[AbstractEventLoop] = None) -> Future:
"""
Execute multiple coroutines concurrently, returning only the results of the first execution.
When one is completed, the execution of other coroutines will be canceled.
"""
loop = loop or asyncio.get_running_loop()
tasks: Set[Task] = set()
result, _future = loop.create_future(), None
def _done_callback(fut: Future) -> None:
nonlocal result, _future
if result.cancelled():
return # nothing to do on onlyfirst cancelled
if _future is None:
_future = fut # record first completed future
cancel_all_task()
if not result.done():
if _future.exception() is None:
result.set_result(_future.result())
else:
result.set_exception(_future.exception())
def cancel_all_task() -> None:
for task in tasks:
task.remove_done_callback(_done_callback)
for task in filter(lambda task: not task.done(), tasks):
task.cancel()
for coro in coros:
task: Task = loop.create_task(coro)
task.add_done_callback(_done_callback)
tasks.add(task)
result.add_done_callback(lambda fut: cancel_all_task())
return result
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment