Spaces:
Running
Running
import functools | |
from typing import Callable, Generator, Iterable, Iterator, Optional, Tuple | |
from pip._vendor.rich.progress import ( | |
BarColumn, | |
DownloadColumn, | |
FileSizeColumn, | |
Progress, | |
ProgressColumn, | |
SpinnerColumn, | |
TextColumn, | |
TimeElapsedColumn, | |
TimeRemainingColumn, | |
TransferSpeedColumn, | |
) | |
from pip._internal.utils.logging import get_indentation | |
DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]] | |
def _rich_progress_bar( | |
iterable: Iterable[bytes], | |
*, | |
bar_type: str, | |
size: int, | |
) -> Generator[bytes, None, None]: | |
assert bar_type == "on", "This should only be used in the default mode." | |
if not size: | |
total = float("inf") | |
columns: Tuple[ProgressColumn, ...] = ( | |
TextColumn("[progress.description]{task.description}"), | |
SpinnerColumn("line", speed=1.5), | |
FileSizeColumn(), | |
TransferSpeedColumn(), | |
TimeElapsedColumn(), | |
) | |
else: | |
total = size | |
columns = ( | |
TextColumn("[progress.description]{task.description}"), | |
BarColumn(), | |
DownloadColumn(), | |
TransferSpeedColumn(), | |
TextColumn("eta"), | |
TimeRemainingColumn(), | |
) | |
progress = Progress(*columns, refresh_per_second=30) | |
task_id = progress.add_task(" " * (get_indentation() + 2), total=total) | |
with progress: | |
for chunk in iterable: | |
yield chunk | |
progress.update(task_id, advance=len(chunk)) | |
def get_download_progress_renderer( | |
*, bar_type: str, size: Optional[int] = None | |
) -> DownloadProgressRenderer: | |
"""Get an object that can be used to render the download progress. | |
Returns a callable, that takes an iterable to "wrap". | |
""" | |
if bar_type == "on": | |
return functools.partial(_rich_progress_bar, bar_type=bar_type, size=size) | |
else: | |
return iter # no-op, when passed an iterator | |