diff --git a/src/pip/_internal/cli/progress_bars.py b/src/pip/_internal/cli/progress_bars.py index 04e9147df73..ffa1964fc7b 100644 --- a/src/pip/_internal/cli/progress_bars.py +++ b/src/pip/_internal/cli/progress_bars.py @@ -1,10 +1,23 @@ +import functools import itertools import sys from signal import SIGINT, default_int_handler, signal -from typing import Any, Callable, Iterator +from typing import Any, Callable, Iterator, Optional, Tuple from pip._vendor.progress.bar import Bar, FillingCirclesBar, IncrementalBar from pip._vendor.progress.spinner import Spinner +from pip._vendor.rich.progress import ( + BarColumn, + DownloadColumn, + FileSizeColumn, + Progress, + ProgressColumn, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, + TimeRemainingColumn, + TransferSpeedColumn, +) from pip._internal.utils.compat import WINDOWS from pip._internal.utils.logging import get_indentation @@ -245,10 +258,64 @@ def update(self) -> None: } -def DownloadProgressProvider( +def _legacy_progress_bar( progress_bar: str, max: Optional[int] ) -> DownloadProgressRenderer: if max is None or max == 0: return BAR_TYPES[progress_bar][1]().iter # type: ignore else: return BAR_TYPES[progress_bar][0](max=max).iter + + +# +# Modern replacement, for our legacy progress bars. +# +def _rich_progress_bar( + iterable: Iterator[bytes], + *, + bar_type: str, + size: int, +) -> Iterator[bytes]: + 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) + elif bar_type == "off": + return iter # no-op, when passed an iterator + else: + return _legacy_progress_bar(bar_type, size) diff --git a/src/pip/_internal/network/download.py b/src/pip/_internal/network/download.py index 47af547d6fc..35bc970e260 100644 --- a/src/pip/_internal/network/download.py +++ b/src/pip/_internal/network/download.py @@ -8,7 +8,7 @@ from pip._vendor.requests.models import CONTENT_CHUNK_SIZE, Response -from pip._internal.cli.progress_bars import DownloadProgressProvider +from pip._internal.cli.progress_bars import get_download_progress_renderer from pip._internal.exceptions import NetworkConnectionError from pip._internal.models.index import PyPI from pip._internal.models.link import Link @@ -65,7 +65,8 @@ def _prepare_download( if not show_progress: return chunks - return DownloadProgressProvider(progress_bar, max=total_length)(chunks) + renderer = get_download_progress_renderer(bar_type=progress_bar, size=total_length) + return renderer(chunks) def sanitize_content_filename(filename: str) -> str: