69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
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
|