Skip to content

Commit

Permalink
use asyncio.TaskGroup to automatically cancel other jobs when one f…
Browse files Browse the repository at this point in the history
…ails
  • Loading branch information
CNSeniorious000 committed Jun 4, 2024
1 parent e01491f commit 88655a6
Showing 1 changed file with 11 additions and 15 deletions.
26 changes: 11 additions & 15 deletions micropip/transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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):
Expand Down

0 comments on commit 88655a6

Please sign in to comment.