| from __future__ import annotations |
|
|
| import functools |
| import sys |
| from collections.abc import Generator, Iterable, Iterator |
| from typing import Callable, Literal, TypeVar |
|
|
| from pip._vendor.rich.progress import ( |
| BarColumn, |
| DownloadColumn, |
| FileSizeColumn, |
| MofNCompleteColumn, |
| Progress, |
| ProgressColumn, |
| SpinnerColumn, |
| TextColumn, |
| TimeElapsedColumn, |
| TimeRemainingColumn, |
| TransferSpeedColumn, |
| ) |
|
|
| from pip._internal.cli.spinners import RateLimiter |
| from pip._internal.req.req_install import InstallRequirement |
| from pip._internal.utils.logging import get_console, get_indentation |
|
|
| T = TypeVar("T") |
| ProgressRenderer = Callable[[Iterable[T]], Iterator[T]] |
| BarType = Literal["on", "off", "raw"] |
|
|
|
|
| def _rich_download_progress_bar( |
| iterable: Iterable[bytes], |
| *, |
| bar_type: BarType, |
| size: int | None, |
| initial_progress: int | None = None, |
| ) -> 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("{task.fields[time_description]}"), |
| TimeRemainingColumn(elapsed_when_finished=True), |
| ) |
|
|
| progress = Progress(*columns, refresh_per_second=5) |
| task_id = progress.add_task( |
| " " * (get_indentation() + 2), total=total, time_description="eta" |
| ) |
| if initial_progress is not None: |
| progress.update(task_id, advance=initial_progress) |
| with progress: |
| for chunk in iterable: |
| yield chunk |
| progress.update(task_id, advance=len(chunk)) |
| progress.update(task_id, time_description="") |
|
|
|
|
| def _rich_install_progress_bar( |
| iterable: Iterable[InstallRequirement], *, total: int |
| ) -> Iterator[InstallRequirement]: |
| columns = ( |
| TextColumn("{task.fields[indent]}"), |
| BarColumn(), |
| MofNCompleteColumn(), |
| TextColumn("{task.description}"), |
| ) |
| console = get_console() |
|
|
| bar = Progress(*columns, refresh_per_second=6, console=console, transient=True) |
| |
| |
| task = bar.add_task("", total=total, indent=" " * get_indentation(), visible=False) |
| with bar: |
| for req in iterable: |
| bar.update(task, description=rf"\[{req.name}]", visible=True) |
| yield req |
| bar.advance(task) |
|
|
|
|
| def _raw_progress_bar( |
| iterable: Iterable[bytes], |
| *, |
| size: int | None, |
| initial_progress: int | None = None, |
| ) -> Generator[bytes, None, None]: |
| def write_progress(current: int, total: int) -> None: |
| sys.stdout.write(f"Progress {current} of {total}\n") |
| sys.stdout.flush() |
|
|
| current = initial_progress or 0 |
| total = size or 0 |
| rate_limiter = RateLimiter(0.25) |
|
|
| write_progress(current, total) |
| for chunk in iterable: |
| current += len(chunk) |
| if rate_limiter.ready() or current == total: |
| write_progress(current, total) |
| rate_limiter.reset() |
| yield chunk |
|
|
|
|
| def get_download_progress_renderer( |
| *, bar_type: BarType, size: int | None = None, initial_progress: int | None = None |
| ) -> ProgressRenderer[bytes]: |
| """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_download_progress_bar, |
| bar_type=bar_type, |
| size=size, |
| initial_progress=initial_progress, |
| ) |
| elif bar_type == "raw": |
| return functools.partial( |
| _raw_progress_bar, |
| size=size, |
| initial_progress=initial_progress, |
| ) |
| else: |
| return iter |
|
|
|
|
| def get_install_progress_renderer( |
| *, bar_type: BarType, total: int |
| ) -> ProgressRenderer[InstallRequirement]: |
| """Get an object that can be used to render the install progress. |
| Returns a callable, that takes an iterable to "wrap". |
| """ |
| if bar_type == "on": |
| return functools.partial(_rich_install_progress_bar, total=total) |
| else: |
| return iter |
|
|
