Skip to content

Commit de44d99

Browse files
authored
Merge pull request #13220 from ichard26/feat/install-bar
Installation progress bar ✨
2 parents ce94515 + 24e364e commit de44d99

File tree

6 files changed

+88
-23
lines changed

6 files changed

+88
-23
lines changed

news/12712.feature.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Display a transient progress bar during package installation.

src/pip/_internal/cli/progress_bars.py

Lines changed: 45 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
import functools
22
import sys
3-
from typing import Callable, Generator, Iterable, Iterator, Optional, Tuple
3+
from typing import Callable, Generator, Iterable, Iterator, Optional, Tuple, TypeVar
44

55
from pip._vendor.rich.progress import (
66
BarColumn,
77
DownloadColumn,
88
FileSizeColumn,
9+
MofNCompleteColumn,
910
Progress,
1011
ProgressColumn,
1112
SpinnerColumn,
@@ -16,12 +17,14 @@
1617
)
1718

1819
from pip._internal.cli.spinners import RateLimiter
19-
from pip._internal.utils.logging import get_indentation
20+
from pip._internal.req.req_install import InstallRequirement
21+
from pip._internal.utils.logging import get_console, get_indentation
2022

21-
DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]]
23+
T = TypeVar("T")
24+
ProgressRenderer = Callable[[Iterable[T]], Iterator[T]]
2225

2326

24-
def _rich_progress_bar(
27+
def _rich_download_progress_bar(
2528
iterable: Iterable[bytes],
2629
*,
2730
bar_type: str,
@@ -57,6 +60,28 @@ def _rich_progress_bar(
5760
progress.update(task_id, advance=len(chunk))
5861

5962

63+
def _rich_install_progress_bar(
64+
iterable: Iterable[InstallRequirement], *, total: int
65+
) -> Iterator[InstallRequirement]:
66+
columns = (
67+
TextColumn("{task.fields[indent]}"),
68+
BarColumn(),
69+
MofNCompleteColumn(),
70+
TextColumn("{task.description}"),
71+
)
72+
console = get_console()
73+
74+
bar = Progress(*columns, refresh_per_second=6, console=console, transient=True)
75+
# Hiding the progress bar at initialization forces a refresh cycle to occur
76+
# until the bar appears, avoiding very short flashes.
77+
task = bar.add_task("", total=total, indent=" " * get_indentation(), visible=False)
78+
with bar:
79+
for req in iterable:
80+
bar.update(task, description=rf"\[{req.name}]", visible=True)
81+
yield req
82+
bar.advance(task)
83+
84+
6085
def _raw_progress_bar(
6186
iterable: Iterable[bytes],
6287
*,
@@ -81,14 +106,28 @@ def write_progress(current: int, total: int) -> None:
81106

82107
def get_download_progress_renderer(
83108
*, bar_type: str, size: Optional[int] = None
84-
) -> DownloadProgressRenderer:
109+
) -> ProgressRenderer[bytes]:
85110
"""Get an object that can be used to render the download progress.
86111
87112
Returns a callable, that takes an iterable to "wrap".
88113
"""
89114
if bar_type == "on":
90-
return functools.partial(_rich_progress_bar, bar_type=bar_type, size=size)
115+
return functools.partial(
116+
_rich_download_progress_bar, bar_type=bar_type, size=size
117+
)
91118
elif bar_type == "raw":
92119
return functools.partial(_raw_progress_bar, size=size)
93120
else:
94121
return iter # no-op, when passed an iterator
122+
123+
124+
def get_install_progress_renderer(
125+
*, bar_type: str, total: int
126+
) -> ProgressRenderer[InstallRequirement]:
127+
"""Get an object that can be used to render the install progress.
128+
Returns a callable, that takes an iterable to "wrap".
129+
"""
130+
if bar_type == "on":
131+
return functools.partial(_rich_install_progress_bar, total=total)
132+
else:
133+
return iter

src/pip/_internal/commands/install.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,7 @@ def run(self, options: Values, args: List[str]) -> int:
464464
warn_script_location=warn_script_location,
465465
use_user_site=options.use_user_site,
466466
pycompile=options.compile,
467+
progress_bar=options.progress_bar,
467468
)
468469

469470
lib_locations = get_lib_location_guesses(

src/pip/_internal/req/__init__.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from dataclasses import dataclass
44
from typing import Generator, List, Optional, Sequence, Tuple
55

6+
from pip._internal.cli.progress_bars import get_install_progress_renderer
67
from pip._internal.utils.logging import indent_log
78

89
from .req_file import parse_requirements
@@ -41,6 +42,7 @@ def install_given_reqs(
4142
warn_script_location: bool,
4243
use_user_site: bool,
4344
pycompile: bool,
45+
progress_bar: str,
4446
) -> List[InstallationResult]:
4547
"""
4648
Install everything in the given list.
@@ -57,8 +59,19 @@ def install_given_reqs(
5759

5860
installed = []
5961

62+
show_progress = logger.isEnabledFor(logging.INFO) and len(to_install) > 1
63+
64+
items = iter(to_install.values())
65+
if show_progress:
66+
renderer = get_install_progress_renderer(
67+
bar_type=progress_bar, total=len(to_install)
68+
)
69+
items = renderer(items)
70+
6071
with indent_log():
61-
for req_name, requirement in to_install.items():
72+
for requirement in items:
73+
req_name = requirement.name
74+
assert req_name is not None
6275
if requirement.should_reinstall:
6376
logger.info("Attempting uninstall: %s", req_name)
6477
with indent_log():

src/pip/_internal/utils/logging.py

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from dataclasses import dataclass
99
from io import TextIOWrapper
1010
from logging import Filter
11-
from typing import Any, ClassVar, Generator, List, Optional, TextIO, Type
11+
from typing import Any, ClassVar, Generator, List, Optional, Type
1212

1313
from pip._vendor.rich.console import (
1414
Console,
@@ -29,6 +29,8 @@
2929
from pip._internal.utils.misc import ensure_dir
3030

3131
_log_state = threading.local()
32+
_stdout_console = None
33+
_stderr_console = None
3234
subprocess_logger = getLogger("pip.subprocessor")
3335

3436

@@ -144,12 +146,21 @@ def on_broken_pipe(self) -> None:
144146
raise BrokenPipeError() from None
145147

146148

149+
def get_console(*, stderr: bool = False) -> Console:
150+
if stderr:
151+
assert _stderr_console is not None, "stderr rich console is missing!"
152+
return _stderr_console
153+
else:
154+
assert _stdout_console is not None, "stdout rich console is missing!"
155+
return _stdout_console
156+
157+
147158
class RichPipStreamHandler(RichHandler):
148159
KEYWORDS: ClassVar[Optional[List[str]]] = []
149160

150-
def __init__(self, stream: Optional[TextIO], no_color: bool) -> None:
161+
def __init__(self, console: Console) -> None:
151162
super().__init__(
152-
console=PipConsole(file=stream, no_color=no_color, soft_wrap=True),
163+
console=console,
153164
show_time=False,
154165
show_level=False,
155166
show_path=False,
@@ -266,17 +277,16 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
266277
vendored_log_level = "WARNING" if level in ["INFO", "ERROR"] else "DEBUG"
267278

268279
# Shorthands for clarity
269-
log_streams = {
270-
"stdout": "ext://sys.stdout",
271-
"stderr": "ext://sys.stderr",
272-
}
273280
handler_classes = {
274281
"stream": "pip._internal.utils.logging.RichPipStreamHandler",
275282
"file": "pip._internal.utils.logging.BetterRotatingFileHandler",
276283
}
277284
handlers = ["console", "console_errors", "console_subprocess"] + (
278285
["user_log"] if include_user_log else []
279286
)
287+
global _stdout_console, stderr_console
288+
_stdout_console = PipConsole(file=sys.stdout, no_color=no_color, soft_wrap=True)
289+
_stderr_console = PipConsole(file=sys.stderr, no_color=no_color, soft_wrap=True)
280290

281291
logging.config.dictConfig(
282292
{
@@ -311,16 +321,14 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
311321
"console": {
312322
"level": level,
313323
"class": handler_classes["stream"],
314-
"no_color": no_color,
315-
"stream": log_streams["stdout"],
324+
"console": _stdout_console,
316325
"filters": ["exclude_subprocess", "exclude_warnings"],
317326
"formatter": "indent",
318327
},
319328
"console_errors": {
320329
"level": "WARNING",
321330
"class": handler_classes["stream"],
322-
"no_color": no_color,
323-
"stream": log_streams["stderr"],
331+
"console": _stderr_console,
324332
"filters": ["exclude_subprocess"],
325333
"formatter": "indent",
326334
},
@@ -329,8 +337,7 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
329337
"console_subprocess": {
330338
"level": level,
331339
"class": handler_classes["stream"],
332-
"stream": log_streams["stderr"],
333-
"no_color": no_color,
340+
"console": _stderr_console,
334341
"filters": ["restrict_to_subprocess"],
335342
"formatter": "indent",
336343
},

tests/unit/test_logging.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from pip._internal.utils.logging import (
1111
BrokenStdoutLoggingError,
1212
IndentingFormatter,
13+
PipConsole,
1314
RichPipStreamHandler,
1415
indent_log,
1516
)
@@ -142,7 +143,8 @@ def test_broken_pipe_in_stderr_flush(self) -> None:
142143
record = self._make_log_record()
143144

144145
with redirect_stderr(StringIO()) as stderr:
145-
handler = RichPipStreamHandler(stream=stderr, no_color=True)
146+
console = PipConsole(file=stderr, no_color=True, soft_wrap=True)
147+
handler = RichPipStreamHandler(console)
146148
with patch("sys.stderr.flush") as mock_flush:
147149
mock_flush.side_effect = BrokenPipeError()
148150
# The emit() call raises no exception.
@@ -165,7 +167,8 @@ def test_broken_pipe_in_stdout_write(self) -> None:
165167
record = self._make_log_record()
166168

167169
with redirect_stdout(StringIO()) as stdout:
168-
handler = RichPipStreamHandler(stream=stdout, no_color=True)
170+
console = PipConsole(file=stdout, no_color=True, soft_wrap=True)
171+
handler = RichPipStreamHandler(console)
169172
with patch("sys.stdout.write") as mock_write:
170173
mock_write.side_effect = BrokenPipeError()
171174
with pytest.raises(BrokenStdoutLoggingError):
@@ -180,7 +183,8 @@ def test_broken_pipe_in_stdout_flush(self) -> None:
180183
record = self._make_log_record()
181184

182185
with redirect_stdout(StringIO()) as stdout:
183-
handler = RichPipStreamHandler(stream=stdout, no_color=True)
186+
console = PipConsole(file=stdout, no_color=True, soft_wrap=True)
187+
handler = RichPipStreamHandler(console)
184188
with patch("sys.stdout.flush") as mock_flush:
185189
mock_flush.side_effect = BrokenPipeError()
186190
with pytest.raises(BrokenStdoutLoggingError):

0 commit comments

Comments
 (0)