Skip to content

Commit

Permalink
Use the progress bar from rich by default
Browse files Browse the repository at this point in the history
Utilise rich's progress bar to present download progress. This has a
subjectively nicer presentation style and should degrade gracefully
without additional effort from our end.
  • Loading branch information
pradyunsg committed Nov 12, 2021
1 parent f2108aa commit 5d07b09
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 4 deletions.
71 changes: 69 additions & 2 deletions src/pip/_internal/cli/progress_bars.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)
5 changes: 3 additions & 2 deletions src/pip/_internal/network/download.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 5d07b09

Please sign in to comment.