322 lines
10 KiB
Python
322 lines
10 KiB
Python
|
import functools
|
||
|
|
import itertools
|
||
|
|
import sys
|
||
|
|
from signal import SIGINT, default_int_handler, signal
|
||
|
|
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
|
||
|
|
from pip._internal.utils.misc import format_size
|
||
|
|
|
||
|
|
try:
|
||
|
|
from pip._vendor import colorama
|
||
|
|
# Lots of different errors can come from this, including SystemError and
|
||
|
|
# ImportError.
|
||
|
|
except Exception:
|
||
|
|
colorama = None
|
||
|
|
|
||
|
|
DownloadProgressRenderer = Callable[[Iterator[bytes]], Iterator[bytes]]
|
||
|
|
|
||
|
|
|
||
|
|
def _select_progress_class(preferred: Bar, fallback: Bar) -> Bar:
|
||
|
|
encoding = getattr(preferred.file, "encoding", None)
|
||
|
|
|
||
|
|
# If we don't know what encoding this file is in, then we'll just assume
|
||
|
|
# that it doesn't support unicode and use the ASCII bar.
|
||
|
|
if not encoding:
|
||
|
|
return fallback
|
||
|
|
|
||
|
|
# Collect all of the possible characters we want to use with the preferred
|
||
|
|
# bar.
|
||
|
|
characters = [
|
||
|
|
getattr(preferred, "empty_fill", ""),
|
||
|
|
getattr(preferred, "fill", ""),
|
||
|
|
]
|
||
|
|
characters += list(getattr(preferred, "phases", []))
|
||
|
|
|
||
|
|
# Try to decode the characters we're using for the bar using the encoding
|
||
|
|
# of the given file, if this works then we'll assume that we can use the
|
||
|
|
# fancier bar and if not we'll fall back to the plaintext bar.
|
||
|
|
try:
|
||
|
|
"".join(characters).encode(encoding)
|
||
|
|
except UnicodeEncodeError:
|
||
|
|
return fallback
|
||
|
|
else:
|
||
|
|
return preferred
|
||
|
|
|
||
|
|
|
||
|
|
_BaseBar: Any = _select_progress_class(IncrementalBar, Bar)
|
||
|
|
|
||
|
|
|
||
|
|
class InterruptibleMixin:
|
||
|
|
"""
|
||
|
|
Helper to ensure that self.finish() gets called on keyboard interrupt.
|
||
|
|
|
||
|
|
This allows downloads to be interrupted without leaving temporary state
|
||
|
|
(like hidden cursors) behind.
|
||
|
|
|
||
|
|
This class is similar to the progress library's existing SigIntMixin
|
||
|
|
helper, but as of version 1.2, that helper has the following problems:
|
||
|
|
|
||
|
|
1. It calls sys.exit().
|
||
|
|
2. It discards the existing SIGINT handler completely.
|
||
|
|
3. It leaves its own handler in place even after an uninterrupted finish,
|
||
|
|
which will have unexpected delayed effects if the user triggers an
|
||
|
|
unrelated keyboard interrupt some time after a progress-displaying
|
||
|
|
download has already completed, for example.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||
|
|
"""
|
||
|
|
Save the original SIGINT handler for later.
|
||
|
|
"""
|
||
|
|
# https://github.com/python/mypy/issues/5887
|
||
|
|
super().__init__(*args, **kwargs) # type: ignore
|
||
|
|
|
||
|
|
self.original_handler = signal(SIGINT, self.handle_sigint)
|
||
|
|
|
||
|
|
# If signal() returns None, the previous handler was not installed from
|
||
|
|
# Python, and we cannot restore it. This probably should not happen,
|
||
|
|
# but if it does, we must restore something sensible instead, at least.
|
||
|
|
# The least bad option should be Python's default SIGINT handler, which
|
||
|
|
# just raises KeyboardInterrupt.
|
||
|
|
if self.original_handler is None:
|
||
|
|
self.original_handler = default_int_handler
|
||
|
|
|
||
|
|
def finish(self) -> None:
|
||
|
|
"""
|
||
|
|
Restore the original SIGINT handler after finishing.
|
||
|
|
|
||
|
|
This should happen regardless of whether the progress display finishes
|
||
|
|
normally, or gets interrupted.
|
||
|
|
"""
|
||
|
|
super().finish() # type: ignore
|
||
|
|
signal(SIGINT, self.original_handler)
|
||
|
|
|
||
|
|
def handle_sigint(self, signum, frame): # type: ignore
|
||
|
|
"""
|
||
|
|
Call self.finish() before delegating to the original SIGINT handler.
|
||
|
|
|
||
|
|
This handler should only be in place while the progress display is
|
||
|
|
active.
|
||
|
|
"""
|
||
|
|
self.finish()
|
||
|
|
self.original_handler(signum, frame)
|
||
|
|
|
||
|
|
|
||
|
|
class SilentBar(Bar):
|
||
|
|
def update(self) -> None:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class BlueEmojiBar(IncrementalBar):
|
||
|
|
|
||
|
|
suffix = "%(percent)d%%"
|
||
|
|
bar_prefix = " "
|
||
|
|
bar_suffix = " "
|
||
|
|
phases = ("\U0001F539", "\U0001F537", "\U0001F535")
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadProgressMixin:
|
||
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||
|
|
# https://github.com/python/mypy/issues/5887
|
||
|
|
super().__init__(*args, **kwargs) # type: ignore
|
||
|
|
self.message: str = (" " * (get_indentation() + 2)) + self.message
|
||
|
|
|
||
|
|
@property
|
||
|
|
def downloaded(self) -> str:
|
||
|
|
return format_size(self.index) # type: ignore
|
||
|
|
|
||
|
|
@property
|
||
|
|
def download_speed(self) -> str:
|
||
|
|
# Avoid zero division errors...
|
||
|
|
if self.avg == 0.0: # type: ignore
|
||
|
|
return "..."
|
||
|
|
return format_size(1 / self.avg) + "/s" # type: ignore
|
||
|
|
|
||
|
|
@property
|
||
|
|
def pretty_eta(self) -> str:
|
||
|
|
if self.eta: # type: ignore
|
||
|
|
return f"eta {self.eta_td}" # type: ignore
|
||
|
|
return ""
|
||
|
|
|
||
|
|
def iter(self, it): # type: ignore
|
||
|
|
for x in it:
|
||
|
|
yield x
|
||
|
|
# B305 is incorrectly raised here
|
||
|
|
# https://github.com/PyCQA/flake8-bugbear/issues/59
|
||
|
|
self.next(len(x)) # noqa: B305
|
||
|
|
self.finish()
|
||
|
|
|
||
|
|
|
||
|
|
class WindowsMixin:
|
||
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||
|
|
# The Windows terminal does not support the hide/show cursor ANSI codes
|
||
|
|
# even with colorama. So we'll ensure that hide_cursor is False on
|
||
|
|
# Windows.
|
||
|
|
# This call needs to go before the super() call, so that hide_cursor
|
||
|
|
# is set in time. The base progress bar class writes the "hide cursor"
|
||
|
|
# code to the terminal in its init, so if we don't set this soon
|
||
|
|
# enough, we get a "hide" with no corresponding "show"...
|
||
|
|
if WINDOWS and self.hide_cursor: # type: ignore
|
||
|
|
self.hide_cursor = False
|
||
|
|
|
||
|
|
# https://github.com/python/mypy/issues/5887
|
||
|
|
super().__init__(*args, **kwargs) # type: ignore
|
||
|
|
|
||
|
|
# Check if we are running on Windows and we have the colorama module,
|
||
|
|
# if we do then wrap our file with it.
|
||
|
|
if WINDOWS and colorama:
|
||
|
|
self.file = colorama.AnsiToWin32(self.file) # type: ignore
|
||
|
|
# The progress code expects to be able to call self.file.isatty()
|
||
|
|
# but the colorama.AnsiToWin32() object doesn't have that, so we'll
|
||
|
|
# add it.
|
||
|
|
self.file.isatty = lambda: self.file.wrapped.isatty()
|
||
|
|
# The progress code expects to be able to call self.file.flush()
|
||
|
|
# but the colorama.AnsiToWin32() object doesn't have that, so we'll
|
||
|
|
# add it.
|
||
|
|
self.file.flush = lambda: self.file.wrapped.flush()
|
||
|
|
|
||
|
|
|
||
|
|
class BaseDownloadProgressBar(WindowsMixin, InterruptibleMixin, DownloadProgressMixin):
|
||
|
|
|
||
|
|
file = sys.stdout
|
||
|
|
message = "%(percent)d%%"
|
||
|
|
suffix = "%(downloaded)s %(download_speed)s %(pretty_eta)s"
|
||
|
|
|
||
|
|
|
||
|
|
class DefaultDownloadProgressBar(BaseDownloadProgressBar, _BaseBar):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadSilentBar(BaseDownloadProgressBar, SilentBar):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadBar(BaseDownloadProgressBar, Bar):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadFillingCirclesBar(BaseDownloadProgressBar, FillingCirclesBar):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadBlueEmojiProgressBar(BaseDownloadProgressBar, BlueEmojiBar):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class DownloadProgressSpinner(
|
||
|
|
WindowsMixin, InterruptibleMixin, DownloadProgressMixin, Spinner
|
||
|
|
):
|
||
|
|
|
||
|
|
file = sys.stdout
|
||
|
|
suffix = "%(downloaded)s %(download_speed)s"
|
||
|
|
|
||
|
|
def next_phase(self) -> str:
|
||
|
|
if not hasattr(self, "_phaser"):
|
||
|
|
self._phaser = itertools.cycle(self.phases)
|
||
|
|
return next(self._phaser)
|
||
|
|
|
||
|
|
def update(self) -> None:
|
||
|
|
message = self.message % self
|
||
|
|
phase = self.next_phase()
|
||
|
|
suffix = self.suffix % self
|
||
|
|
line = "".join(
|
||
|
|
[
|
||
|
|
message,
|
||
|
|
" " if message else "",
|
||
|
|
phase,
|
||
|
|
" " if suffix else "",
|
||
|
|
suffix,
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
self.writeln(line)
|
||
|
|
|
||
|
|
|
||
|
|
BAR_TYPES = {
|
||
|
|
"off": (DownloadSilentBar, DownloadSilentBar),
|
||
|
|
"on": (DefaultDownloadProgressBar, DownloadProgressSpinner),
|
||
|
|
"ascii": (DownloadBar, DownloadProgressSpinner),
|
||
|
|
"pretty": (DownloadFillingCirclesBar, DownloadProgressSpinner),
|
||
|
|
"emoji": (DownloadBlueEmojiProgressBar, DownloadProgressSpinner),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
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)
|