mirror of https://github.com/pypa/pip
375 lines
12 KiB
Python
375 lines
12 KiB
Python
import functools
|
|
from logging import Logger
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Generator,
|
|
Iterable,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
from pip._vendor.rich.console import (
|
|
Console,
|
|
ConsoleOptions,
|
|
RenderableType,
|
|
RenderResult,
|
|
)
|
|
from pip._vendor.rich.progress import (
|
|
BarColumn,
|
|
DownloadColumn,
|
|
FileSizeColumn,
|
|
Progress,
|
|
ProgressColumn,
|
|
SpinnerColumn,
|
|
Task,
|
|
TaskID,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
TimeRemainingColumn,
|
|
TransferSpeedColumn,
|
|
)
|
|
from pip._vendor.rich.progress_bar import ProgressBar
|
|
from pip._vendor.rich.segment import Segment
|
|
from pip._vendor.rich.text import Text
|
|
|
|
from pip._internal.utils.logging import get_indentation
|
|
|
|
DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]]
|
|
|
|
|
|
class RenderableLine:
|
|
"""
|
|
A wrapper for a single row, renderable by `Console` methods.
|
|
"""
|
|
|
|
def __init__(self, line_items: List[Union[Text, ProgressBar]]):
|
|
self.line_items = line_items
|
|
|
|
def __rich_console__(
|
|
self, console: Console, options: ConsoleOptions
|
|
) -> RenderResult:
|
|
for line_item in self.line_items:
|
|
segments = [
|
|
seg
|
|
for seg in line_item.__rich_console__(console, options)
|
|
if isinstance(seg, Segment) and seg.text != "\n"
|
|
]
|
|
yield from segments
|
|
|
|
|
|
class RenderableLines:
|
|
"""
|
|
A wrapper for multiple rows, renderable by `Console` methods.
|
|
"""
|
|
|
|
def __init__(self, lines: Iterable[RenderableLine]):
|
|
self.lines = lines
|
|
|
|
def __rich_console__(
|
|
self, console: Console, options: ConsoleOptions
|
|
) -> RenderResult:
|
|
for idx, line in enumerate(self.lines):
|
|
if idx != 0:
|
|
yield Segment.line()
|
|
yield from line.__rich_console__(console, options)
|
|
yield Segment.line()
|
|
|
|
|
|
class PipProgress(Progress):
|
|
"""
|
|
Custom Progress bar for sequential downloads.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
refresh_per_second: int,
|
|
progress_disabled: bool = False,
|
|
logger: Optional[Logger] = None,
|
|
) -> None:
|
|
super().__init__(refresh_per_second=refresh_per_second)
|
|
self.progress_disabled = progress_disabled
|
|
self.log_download_description = True
|
|
self.logger = logger
|
|
|
|
@classmethod
|
|
def get_default_columns(cls) -> Tuple[ProgressColumn, ...]:
|
|
"""
|
|
Get the default columns to use for the progress bar when the size of
|
|
the file is known.
|
|
"""
|
|
return (
|
|
TextColumn(" " * (get_indentation() + 3)),
|
|
BarColumn(),
|
|
DownloadColumn(),
|
|
TransferSpeedColumn(),
|
|
TextColumn("eta"),
|
|
TimeRemainingColumn(),
|
|
)
|
|
|
|
@classmethod
|
|
def get_indefinite_columns(cls) -> Tuple[ProgressColumn, ...]:
|
|
"""
|
|
Get the columns to use for the progress bar when the size of the file
|
|
is unknown.
|
|
"""
|
|
return (
|
|
TextColumn(" " * (get_indentation() + 3)),
|
|
SpinnerColumn("line", speed=1.5),
|
|
FileSizeColumn(),
|
|
TransferSpeedColumn(),
|
|
TimeElapsedColumn(),
|
|
)
|
|
|
|
def get_renderable(self) -> RenderableType:
|
|
"""
|
|
Get the renderable representation of the progress of all tasks
|
|
"""
|
|
renderables: List[RenderableLine] = []
|
|
for task in self.tasks:
|
|
if not task.visible:
|
|
continue
|
|
task_renderable = [x for x in self.make_task_group(task) if x is not None]
|
|
renderables.extend(task_renderable)
|
|
return RenderableLines(renderables)
|
|
|
|
def make_task_group(self, task: Task) -> Iterable[Optional[RenderableLine]]:
|
|
"""
|
|
Create a representation for a task, i.e. it's progress bar.
|
|
|
|
Parameters:
|
|
- task (Task): The task for which to generate the representation.
|
|
|
|
Returns:
|
|
- Optional[Group]: text representation of a Progress Column,
|
|
"""
|
|
|
|
hide_progress = task.fields["hide_progress"]
|
|
if self.progress_disabled or hide_progress:
|
|
return (None,)
|
|
columns = (
|
|
self.columns if task.total is not None else self.get_indefinite_columns()
|
|
)
|
|
progress_row = self.make_task_row(columns, task)
|
|
return (progress_row,)
|
|
|
|
def make_task_row(
|
|
self, columns: Tuple[Union[str, ProgressColumn], ...], task: Task
|
|
) -> RenderableLine:
|
|
"""
|
|
Format the columns of a single row for the given task.
|
|
|
|
"""
|
|
row_values = [
|
|
(column.format(task=task) if isinstance(column, str) else column(task))
|
|
for column in columns
|
|
]
|
|
row = self.merge_text_objects(row_values)
|
|
return RenderableLine(row)
|
|
|
|
def merge_text_objects(
|
|
self, row: List[RenderableType]
|
|
) -> List[Union[Text, ProgressBar]]:
|
|
"""
|
|
Merge adjacent Text objects in the given row into a single Text object.
|
|
"""
|
|
merged_row: List[Union[Text, ProgressBar]] = []
|
|
markup_to_merge: List[str] = []
|
|
for item in row:
|
|
if isinstance(item, ProgressBar):
|
|
if markup_to_merge:
|
|
merged_text = Text.from_markup(" ".join(markup_to_merge))
|
|
merged_row.append(merged_text)
|
|
merged_row.append(item)
|
|
markup_to_merge = [" "]
|
|
elif isinstance(item, Text):
|
|
markup_to_merge.append(item.markup)
|
|
if markup_to_merge:
|
|
merged_markup = " ".join(markup_to_merge)
|
|
merged_row.append(Text.from_markup(merged_markup))
|
|
return merged_row
|
|
|
|
def add_task(
|
|
self,
|
|
description: str,
|
|
start: bool = True,
|
|
total: Optional[float] = 100.0,
|
|
completed: int = 0,
|
|
visible: bool = True,
|
|
**fields: Any,
|
|
) -> TaskID:
|
|
"""
|
|
Reimplementation of Progress.add_task with description logging
|
|
"""
|
|
if visible and self.log_download_description and self.logger:
|
|
indentation = " " * get_indentation()
|
|
log_statement = f"{indentation}{description}"
|
|
self.logger.info(log_statement)
|
|
return super().add_task(
|
|
description=description,
|
|
total=total,
|
|
visible=visible,
|
|
completed=completed,
|
|
**fields,
|
|
)
|
|
|
|
|
|
class PipParallelProgress(PipProgress):
|
|
def __init__(self, refresh_per_second: int, progress_disabled: bool = True):
|
|
super().__init__(
|
|
refresh_per_second=refresh_per_second, progress_disabled=progress_disabled
|
|
)
|
|
# Overrides behaviour of logging description on add_task from PipProgress
|
|
self.log_download_description = False
|
|
|
|
@classmethod
|
|
def get_description_columns(cls) -> Tuple[ProgressColumn, ...]:
|
|
"""
|
|
Get the columns to use for the log message, i.e. the task description
|
|
"""
|
|
# These columns will be the "Downloading"/"Using cached" message
|
|
# This message needs to be columns because,logging this message interferes
|
|
# with parallel progress bars, and if we want the message to remain next
|
|
# to the progress bar even when there are multiple tasks, then it needs
|
|
# to be a part of the progress bar
|
|
indentation = get_indentation()
|
|
if indentation:
|
|
return (
|
|
TextColumn(" " * get_indentation()),
|
|
TextColumn("{task.description}"),
|
|
)
|
|
return (TextColumn("{task.description}"),)
|
|
|
|
def make_task_group(self, task: Task) -> Iterable[Optional[RenderableLine]]:
|
|
"""
|
|
Create a representation for a task, including both the description row
|
|
and the progress row.
|
|
|
|
Parameters:
|
|
- task (Task): The task for which to generate the representation.
|
|
|
|
Returns:
|
|
- Iterable[Optional[RenderableLine]]: An Iterable containing the
|
|
description and progress rows,
|
|
"""
|
|
progress_row = super().make_task_group(task)
|
|
|
|
description_row = self.make_task_row(self.get_description_columns(), task)
|
|
return (description_row, *progress_row)
|
|
|
|
def sort_tasks(self) -> None:
|
|
"""
|
|
Sort tasks
|
|
Remove completed tasks and print them
|
|
"""
|
|
# Removal of completed tasks reduces the number of tasks to be rendered
|
|
tasks = []
|
|
for task_id in self._tasks:
|
|
task = self._tasks[task_id]
|
|
if task.finished and len(self._tasks) > 1:
|
|
# Remove and log the finished task if there are too many active
|
|
# tasks to reduce the number of things to be rendered
|
|
# If there are too many active tasks on screen rich renders the
|
|
# overflow as a ... at the bottom of the screen which makes it
|
|
# difficult for a user to see whats happening
|
|
# If we remove every task on completion, it adds an extra newline
|
|
# for sequential downloads due to self.live on __exit__
|
|
if task.visible:
|
|
task_group = [
|
|
x for x in self.make_task_group(task) if x is not None
|
|
]
|
|
self.console.print(RenderableLines(task_group))
|
|
else:
|
|
tasks.append((task_id, self._tasks[task_id]))
|
|
# Sorting by finished ensures that all active downloads remain together
|
|
tasks = sorted(tasks, key=lambda x: not x[1].finished)
|
|
self._tasks = dict(tasks)
|
|
|
|
def update(
|
|
self,
|
|
task_id: TaskID,
|
|
*,
|
|
total: Optional[float] = None,
|
|
completed: Optional[float] = None,
|
|
advance: Optional[float] = None,
|
|
description: Optional[str] = None,
|
|
visible: Optional[bool] = None,
|
|
refresh: bool = False,
|
|
**fields: Any,
|
|
) -> None:
|
|
"""
|
|
A copy of Progress' implementation of update, with sorting of
|
|
self.tasks when a task is completed
|
|
"""
|
|
with self._lock:
|
|
task = self._tasks[task_id]
|
|
initial_finish_time = task.finished_time
|
|
super().update(
|
|
task_id,
|
|
total=total,
|
|
completed=completed,
|
|
advance=advance,
|
|
description=description,
|
|
visible=visible,
|
|
refresh=refresh,
|
|
**fields,
|
|
)
|
|
# If at the start of the update, the finish time is None and after
|
|
# calling super.update the finish time is not None, it means the
|
|
# task was just finished
|
|
task = self._tasks[task_id]
|
|
if initial_finish_time is None and task.finished_time is not None:
|
|
self.sort_tasks()
|
|
|
|
|
|
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
|