diff --git a/asv.conf.json b/asv.conf.json new file mode 100644 index 0000000000..4e91ee493a --- /dev/null +++ b/asv.conf.json @@ -0,0 +1,23 @@ +{ + "version": 1, + "project": "pip", + "project_url": "https://github.com/pypa/pip", + "repo": ".", + "branches": ["main", "pr-14026"], + "pythons": ["3.11"], + "environment_type": "virtualenv", + "benchmark_dir": "benchmarks", + "results_dir": ".asv/results", + "html_dir": ".asv/html", + "build_command": [ + "python -m ensurepip --upgrade", + "python -m pip wheel -w {build_cache_dir} {build_dir}" + ], + "install_command": [ + "python -m ensurepip --upgrade", + "in-dir={env_dir} python -m pip install {wheel_file} --force-reinstall --disable-pip-version-check --no-index" + ], + "uninstall_command": [], + "matrix": {}, + "build_cache_size": 2 +} diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000000..db29b720a3 --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""ASV benchmark package for pip.""" diff --git a/benchmarks/spinner.py b/benchmarks/spinner.py new file mode 100644 index 0000000000..bfa8ed2a26 --- /dev/null +++ b/benchmarks/spinner.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import io +from dataclasses import dataclass + +from pip._internal.cli import spinners +from pip._internal.utils import logging as pip_logging + + +class CountingTTY(io.TextIOBase): + encoding = "utf-8" + + def __init__(self) -> None: + self.write_calls = 0 + self.flush_calls = 0 + self.bytes_written = 0 + + def isatty(self) -> bool: + return True + + def write(self, text: str) -> int: + self.write_calls += 1 + self.bytes_written += len(text.encode(self.encoding, "replace")) + return len(text) + + def flush(self) -> None: + self.flush_calls += 1 + + +class FakeClock: + def __init__(self) -> None: + self.value = 0.0 + + def now(self) -> float: + return self.value + + def advance(self, seconds: float) -> None: + self.value += seconds + + +@dataclass +class RunResult: + write_calls: int + flush_calls: int + bytes_written: int + + +class TimeSpinnerHotLoop: + """ + Benchmark the checked-out revision's actual ``open_spinner()`` path. + + The workload uses a virtual clock so the benchmark measures spinner + overhead rather than sleeping. This is the path used by + ``runner_with_spinner_message()`` for interactive subprocess status. + """ + + params = ([1, 10, 50],) + param_names = ["packages"] + timeout = 300 + + def _run(self, packages: int) -> RunResult: + stream = CountingTTY() + clock = FakeClock() + + original_spinner_stdout = spinners.sys.stdout + original_logging_stdout = pip_logging.sys.stdout + original_spinner_level = spinners.logger.level + original_console = getattr(pip_logging, "_stdout_console", None) + original_time = spinners.time.time + + try: + spinners.sys.stdout = stream + pip_logging.sys.stdout = stream + spinners.logger.setLevel(spinners.logging.INFO) + spinners.time.time = clock.now + if hasattr(pip_logging, "_stdout_console"): + pip_logging._stdout_console = None + + for package_index in range(packages): + with spinners.open_spinner( + f"Building wheel for package {package_index + 1}/{packages}" + ) as spinner: + for _ in range(50_000): + spinner.spin() + clock.advance(0.0025) + finally: + spinners.sys.stdout = original_spinner_stdout + pip_logging.sys.stdout = original_logging_stdout + spinners.logger.setLevel(original_spinner_level) + spinners.time.time = original_time + if hasattr(pip_logging, "_stdout_console"): + pip_logging._stdout_console = original_console + + return RunResult( + write_calls=stream.write_calls, + flush_calls=stream.flush_calls, + bytes_written=stream.bytes_written, + ) + + def time_spinner_hot_loop(self, packages: int) -> None: + self._run(packages) + + def track_write_calls(self, packages: int) -> int: + return self._run(packages).write_calls + + def track_flush_calls(self, packages: int) -> int: + return self._run(packages).flush_calls + + def track_bytes_written(self, packages: int) -> int: + return self._run(packages).bytes_written diff --git a/news/14028.feature.rst b/news/14028.feature.rst new file mode 100644 index 0000000000..9fe3ff891f --- /dev/null +++ b/news/14028.feature.rst @@ -0,0 +1 @@ +Use Rich's own spinner for better performance. diff --git a/pyproject.toml b/pyproject.toml index 7ae46c0c64..f6a22d7631 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,8 +103,14 @@ docs = [ "sphinx-issues", ] +benchmark = [ + "asv", +] + [tool.flit.sdist] include = [ + "asv.conf.json", + "benchmarks/**/*.py", "NEWS.rst", "SECURITY.md", "AI_POLICY.md", diff --git a/src/pip/_internal/cli/spinners.py b/src/pip/_internal/cli/spinners.py index 58aad2853d..b30d806d5c 100644 --- a/src/pip/_internal/cli/spinners.py +++ b/src/pip/_internal/cli/spinners.py @@ -8,15 +8,8 @@ from collections.abc import Generator from typing import IO, Final -from pip._vendor.rich.console import ( - Console, - ConsoleOptions, - RenderableType, - RenderResult, -) -from pip._vendor.rich.live import Live -from pip._vendor.rich.measure import Measurement -from pip._vendor.rich.text import Text +from pip._vendor.rich.console import Console +from pip._vendor.rich.status import Status from pip._internal.utils.compat import WINDOWS from pip._internal.utils.logging import get_console, get_indentation @@ -128,6 +121,33 @@ def reset(self) -> None: self._last_update = time.time() +class RichStatusSpinner: + def __init__(self, message: str, console: Console | None = None) -> None: + self._message = message + self._label = " " * get_indentation() + message + self._console = console or get_console() + self._status: Status | None = None + if getattr(self._console.file, "isatty", lambda: False)(): + self._status = Status( + f"{self._label} ...", console=self._console, spinner="line" + ) + self._status.__enter__() + self._finished = False + + def finish(self, final_status: str) -> None: + if self._finished: + return + if self._status is not None: + self._status.update(f"{self._label} ... {final_status}") + self._status.__exit__(None, None, None) + self._console.file.write(f"{self._label} ... {final_status}\n") + self._console.file.flush() + else: + self._console.file.write(f"{self._message} ... {final_status}") + self._console.file.flush() + self._finished = True + + @contextlib.contextmanager def open_spinner(message: str) -> Generator[SpinnerInterface, None, None]: # Interactive spinner goes directly to sys.stdout rather than being routed @@ -152,44 +172,6 @@ def open_spinner(message: str) -> Generator[SpinnerInterface, None, None]: spinner.finish("done") -class _PipRichSpinner: - """ - Custom rich spinner that matches the style of the legacy spinners. - - (*) Updates will be handled in a background thread by a rich live panel - which will call render() automatically at the appropriate time. - """ - - def __init__(self, label: str) -> None: - self.label = label - self._spin_cycle = itertools.cycle(SPINNER_CHARS) - self._spinner_text = "" - self._finished = False - self._indent = get_indentation() * " " - - def __rich_console__( - self, console: Console, options: ConsoleOptions - ) -> RenderResult: - yield self.render() - - def __rich_measure__( - self, console: Console, options: ConsoleOptions - ) -> Measurement: - text = self.render() - return Measurement.get(console, options, text) - - def render(self) -> RenderableType: - if not self._finished: - self._spinner_text = next(self._spin_cycle) - - return Text.assemble(self._indent, self.label, " ... ", self._spinner_text) - - def finish(self, status: str) -> None: - """Stop spinning and set a final status message.""" - self._spinner_text = status - self._finished = True - - @contextlib.contextmanager def open_rich_spinner(label: str, console: Console | None = None) -> Generator[None]: if not logger.isEnabledFor(logging.INFO): @@ -197,19 +179,17 @@ def open_rich_spinner(label: str, console: Console | None = None) -> Generator[N yield return - console = console or get_console() - spinner = _PipRichSpinner(label) - with Live(spinner, refresh_per_second=SPINS_PER_SECOND, console=console): - try: - yield - except KeyboardInterrupt: - spinner.finish("canceled") - raise - except Exception: - spinner.finish("error") - raise - else: - spinner.finish("done") + spinner = RichStatusSpinner(label, console=console) + try: + yield + except KeyboardInterrupt: + spinner.finish("canceled") + raise + except Exception: + spinner.finish("error") + raise + else: + spinner.finish("done") HIDE_CURSOR = "\x1b[?25l" diff --git a/src/pip/_internal/utils/logging.py b/src/pip/_internal/utils/logging.py index e067703580..c5f5e473b2 100644 --- a/src/pip/_internal/utils/logging.py +++ b/src/pip/_internal/utils/logging.py @@ -32,8 +32,9 @@ from pip._internal.utils.misc import StreamWrapper, ensure_dir _log_state = threading.local() -_stdout_console = None -_stderr_console = None +_stdout_console: Console | None = None +_stderr_console: Console | None = None + subprocess_logger = getLogger("pip.subprocessor") @@ -190,6 +191,23 @@ def get_console(*, stderr: bool = False) -> Console: return _stdout_console +def get_console_or_create() -> Console: + """Return the stdout console, creating one if logging is not configured yet.""" + global _stdout_console + if _stdout_console is None: + _stdout_console = PipConsole( + file=sys.stdout, + no_color=( + "--no-color" in sys.argv + or os.getenv("PIP_NO_COLOR") == "1" + or "NO_COLOR" in os.environ + ), + soft_wrap=True, + force_terminal=True, + ) + return _stdout_console + + class RichPipStreamHandler(RichHandler): KEYWORDS: ClassVar[list[str] | None] = [] diff --git a/tests/functional/test_wheel.py b/tests/functional/test_wheel.py index 2c3e2e6eb7..583b401d0d 100644 --- a/tests/functional/test_wheel.py +++ b/tests/functional/test_wheel.py @@ -192,7 +192,21 @@ def test_pip_wheel_builds_when_no_binary_set( data.find_links, "simple==3.0", ) - assert "Building wheel for simple" in str(res), str(res) + expected = ( + rf"Looking in links: {re.escape(data.find_links)}\n" + rf"Processing {re.escape(os.fspath(data.packages / 'simple-3.0.tar.gz'))}\n" + r" Preparing metadata \(pyproject\.toml\): started\n" + r" Preparing metadata \(pyproject\.toml\): finished with status 'done'\n" + r"Building wheels for collected packages: simple\n" + r" Building wheel for simple \(pyproject\.toml\): started\n" + r" Building wheel for simple \(pyproject\.toml\): " + r"finished with status 'done'\n" + rf" Created wheel for simple: filename=simple-3.0-py{pyversion[0]}-" + r"none-any\.whl size=\d+ sha256=[a-f0-9]{64}\n" + r" Stored in directory: .+\n" + r"Successfully built simple\n" + ) + assert re.fullmatch(expected, res.stdout), str(res) @pytest.mark.skipif("sys.platform == 'win32'") diff --git a/tests/unit/test_cli_spinners.py b/tests/unit/test_cli_spinners.py index 43736dfe8d..236fc38b8a 100644 --- a/tests/unit/test_cli_spinners.py +++ b/tests/unit/test_cli_spinners.py @@ -26,6 +26,49 @@ def patch_logger_level(level: int) -> Generator[None]: class TestRichSpinner: + def test_interactive_output_keeps_finished_line( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + class TTY(StringIO): + def isatty(self) -> bool: + return True + + monkeypatch.setattr(spinners, "get_indentation", lambda: 2) + stream = TTY() + with patch_logger_level(logging.INFO): + with open_rich_spinner( + "working", Console(file=stream, force_terminal=True) + ): + pass + + output = stream.getvalue() + assert "\x1b[2K working ... done\n" in output + + def test_non_interactive_output( + self, + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr( + spinners, + "sys", + Mock(stdout=Mock(isatty=Mock(return_value=False))), + ) + caplog.set_level(logging.INFO, logger=spinners.logger.name) + + with patch_logger_level(logging.INFO): + with spinners.open_spinner("working"): + pass + + assert [ + record.getMessage() + for record in caplog.records + if record.name == spinners.logger.name + ] == [ + "working: started", + "working: finished with status 'done'", + ] + @pytest.mark.parametrize( "status, func", [ diff --git a/tools/benchmark_spinner_pty.py b/tools/benchmark_spinner_pty.py new file mode 100644 index 0000000000..2313b7b7ef --- /dev/null +++ b/tools/benchmark_spinner_pty.py @@ -0,0 +1,109 @@ +#!/usr/bin/env -S uv run --script +# /// script +# dependencies = ["pyte"] +# /// +"""Exercise pip's spinner through a pseudo-terminal. + +Usage:: + + uv run tools/benchmark_spinner_pty.py /path/to/python --label main + uv run tools/benchmark_spinner_pty.py /path/to/python --label pr + +The supplied Python executable must have pip installed. Run the command once +for each revision, then compare the reported terminal transcript metrics. +This is intentionally separate from the ASV benchmark: PTYs are not +available on every platform, and this measures terminal behavior rather than +just the spinner hot loop. +""" + +from __future__ import annotations + +import argparse +import json +import os +import pty +import select +import subprocess +import time +from pathlib import Path + +import pyte + +CHILD = r""" +import logging +import sys +import time + +from pip._internal.cli import spinners + +spinners.logger.setLevel(logging.INFO) +with spinners.open_spinner("Building wheel") as spinner: + end = time.monotonic() + float(sys.argv[1]) + while time.monotonic() < end: + spinner.spin() + time.sleep(0.0025) +""" + + +def run(python: Path, seconds: float) -> bytes: + master, slave = pty.openpty() + try: + process = subprocess.Popen( + [str(python), "-c", CHILD, str(seconds)], + stdin=slave, + stdout=slave, + stderr=slave, + env={**os.environ, "TERM": "xterm-256color", "COLUMNS": "100"}, + ) + finally: + os.close(slave) + + chunks: list[bytes] = [] + try: + while select.select([master], [], [], 20)[0]: + try: + chunk = os.read(master, 65536) + except OSError: # EIO is the normal PTY EOF on Linux. + break + if not chunk: + break + chunks.append(chunk) + finally: + os.close(master) + if process.wait() != 0: + raise RuntimeError(f"{python} exited with status {process.returncode}") + return b"".join(chunks) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("python", type=Path) + parser.add_argument("--label", default="run") + parser.add_argument("--seconds", type=float, default=5.0) + args = parser.parse_args() + + started = time.monotonic() + transcript = run(args.python, args.seconds) + elapsed = time.monotonic() - started + + screen = pyte.Screen(100, 40) + pyte.ByteStream(screen).feed(transcript) + display = "\n".join(screen.display) + result = { + "label": args.label, + "python": str(args.python), + "seconds": args.seconds, + "elapsed_seconds": round(elapsed, 3), + "bytes_written": len(transcript), + "erase_sequences": transcript.count(b"[2K"), + "cursor_up_sequences": transcript.count(b"[1A"), + "done_line_kept": "... done" in display, + } + print(json.dumps(result, indent=2)) + print("\nFinal terminal screen:\n" + display.rstrip()) + + +if __name__ == "__main__": + if os.name != "posix": + raise SystemExit("This benchmark requires a POSIX pseudo-terminal.") + main()