Add Progress To asyncio Context With TaskGroup and Worker

[YTKme]

Hello,

I am trying to figure how to use rich.progress within and along with other context such as TaskGroup. But struggling to get it integrated and working, hoping someone can share some insight?

I have the below section of code (full code here) to download a list of file provided by source_list.

On a high level, it has a main function (_retrieve_recursive) with a context of TaskGroup, puts the source_list (and potential extra) into a queue, creates tasks, and call worker function (_download_query) to process the queue to download files.

Problem

My current problem is I cannot get a basic progress to correctly show the download progress, and was wondering what is the proper way to achieve it?

• I tried wrapping the with Progress() as progress context before the group.create_task, and passing the progress object down to the _download_query function and update within the _download_query function. But no progress is showing up.
• If I wrap the with Progress() as progress context before async with (session.get...), it gave me an error stating Only one live display may be active at once.
• I tried creating a separate object completely with progress = Progress(...), and using that as a context for async with (session.get...), it seem to have created the progress and shows, but it shows incorrectly and some task doesn't show completion before the worker task ended, so it shows the progress bar stopped at below 100%.

Main Function

async def _retrieve_recursive(
    self,
    source_list: list[str],
    destination_list: list[PathLike],
    download_queue: Queue,
    maximum_connection: int,
    session: ClientSession,
    quiet: bool,
):
    source_queue = Queue()
    destination_queue = Queue()

    # Omitted...
    ...

    async with TaskGroup() as group:
        connection_count = min(download_queue.qsize(), maximum_connection)
        _ = [
            group.create_task(
                self._download_query(
                    download_queue=download_queue,
                    session=session,
                )
            ) for _ in range(connection_count)
        ]

        for _ in range(connection_count):
            await download_queue.put(None)

Worker Function

async def _download_query(
    self,
    download_queue: Queue,
    session: ClientSession,
):
    while True:
        download = await download_queue.get()

        if download is None:
            break

        remote_path = RemotePath(path=download, api_key=self._api_key)

        try:
            Path(unquote(remote_path.location)).parent.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            raise(f'Operating System Error: {e}')

        async with (
            session.get(url=str(remote_path), headers=self._header) as response,
            aiofiles.open(unquote(remote_path.location), 'wb') as file,
        ):
            async for chunk, _ in response.content.iter_chunks():
                await file.write(chunk)