From 88655a6e6782b3205403e1c90802a59880f68473 Mon Sep 17 00:00:00 2001 From: Muspi Merol Date: Tue, 4 Jun 2024 12:13:39 +0800 Subject: [PATCH] use `asyncio.TaskGroup` to automatically cancel other jobs when one fails --- micropip/transaction.py | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/micropip/transaction.py b/micropip/transaction.py index 6686157..e2b86bb 100644 --- a/micropip/transaction.py +++ b/micropip/transaction.py @@ -37,6 +37,8 @@ class Transaction: verbose: bool | int = False + tg: asyncio.TaskGroup | None = None + def __post_init__(self): # If index_urls is None, pyodide-lock.json have to be searched first. # TODO: when PyPI starts to support hosting WASM wheels, this might be deleted. @@ -48,21 +50,15 @@ async def gather_requirements( self, requirements: list[str] | list[Requirement], ) -> None: - requirement_promises = [] - for requirement in requirements: - requirement_promises.append(self.add_requirement(requirement)) - - futures: list[asyncio.Future] = [] - try: - for coro in requirement_promises: - futures.append(asyncio.ensure_future(coro)) - await asyncio.gather(*futures) - except ValueError: - if not self.keep_going: - for future in futures: - if not future.done(): - future.cancel() - raise + if self.tg: + for requirement in requirements: + self.tg.create_task(self.add_requirement(requirement)) + else: + async with asyncio.TaskGroup() as tg: + self.tg = tg # only one task group from the top level + for requirement in requirements: + tg.create_task(self.add_requirement(requirement)) + self.tg = None async def add_requirement(self, req: str | Requirement) -> None: if isinstance(req, Requirement):