Add asyncio.gather support to tqdm.notebook #1665
Description
- I have marked all applicable categories:
- documentation request (i.e. "X is missing from the documentation." If instead I want to ask "how to use X?" I understand StackOverflow#tqdm is more appropriate)
- new feature request
- I have visited the source website, and in particular
read the known issues - I have searched through the issue tracker for duplicates
- I have mentioned version numbers, operating system and
environment, where applicable:
4.67.0 3.11.11 | packaged by conda-forge | (main, Dec 5 2024, 14:17:24) [GCC 13.3.0] linux
Current behavior
When using asyncio.gather in the notebook it is very tricky to use tqdm progress bar to track completed tasks.
from typing import Awaitable, Any
from tqdm.notebook import tqdm
import asyncio
async def gather_with_progress(awaitables: list[Awaitable], **tqdm_kwargs) -> list[Any]:
"""
asyncio.gather with a progress bar.
beuatiful progress bar.
try it out.
"""
async def wrap(i, aw):
result = await aw
return i, result
tasks = [wrap(i, task) for i, task in enumerate(awaitables)]
results = [None] * len(tasks)
progress_bar = tqdm(total=len(tasks), **tqdm_kwargs)
for fut in asyncio.as_completed(tasks):
i, result = await fut
results[i] = result
progress_bar.update(1)
progress_bar.close()
return results
async def work(i):
await asyncio.sleep(i)
return i
tasks = [work(i) for i in range(5)]
results = await gather_with_progress(tasks, ncols=400, leave=True)
print(results)
Expected behavior
Native notebook widget bar that understands coroutines:
from typing import Awaitable, Any
from tqdm.notebook import gather
import asyncio
async def work(i):
await asyncio.sleep(i)
return i
tasks = [work(i) for i in range(5)]
results = await gather(tasks)
print(results)
Proposal
Add as_completed, tarange and gather functions like in tqdm_asyncio to tqdm_notebook.
Context
Currently, only tqdm.asyncio supports gather, but it prints to stderr and lacks notebook widgets (broken for nested loops)
GitHub. This feature would unify the API and improve UX for Jupyter users running concurrent tasks.
I am willing to implement it and contribute to the package.