diff --git a/CHANGELOG.md b/CHANGELOG.md index 522f0690..1d059dd7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ All versions prior to 0.0.9 are untracked. ## [Unreleased] +### Changed + +* CLI: `pip-audit`'s progress spinner has been refactored to make it + faster and more responsive + ([#283](https://github.com/trailofbits/pip-audit/pull/283)) + ## [2.3.1] - 2022-05-24 ### Fixed diff --git a/pip_audit/_cli.py b/pip_audit/_cli.py index daa01e88..e66a2772 100644 --- a/pip_audit/_cli.py +++ b/pip_audit/_cli.py @@ -360,7 +360,7 @@ def audit() -> None: with ExitStack() as stack: actors = [] if args.progress_spinner: - actors.append(AuditSpinner()) + actors.append(AuditSpinner("Collecting inputs")) state = stack.enter_context(AuditState(members=actors)) source: DependencySource diff --git a/pip_audit/_state.py b/pip_audit/_state.py index d16aef57..ba6e2162 100644 --- a/pip_audit/_state.py +++ b/pip_audit/_state.py @@ -4,13 +4,11 @@ """ import logging -import os from abc import ABC, abstractmethod from logging.handlers import MemoryHandler -from typing import Any, Dict, List, Sequence +from typing import Any, List, Sequence -from progress import SHOW_CURSOR -from progress.spinner import Spinner as BaseSpinner +from rich.console import Console class AuditState: @@ -99,24 +97,20 @@ def finalize(self) -> None: raise NotImplementedError # pragma: no cover -class AuditSpinner(_StateActor, BaseSpinner): # pragma: no cover +class AuditSpinner(_StateActor): # pragma: no cover """ - A progress spinner for the `pip-audit` CLI, specialized from `BaseSpinner`. - - This spinner is also written as a `AuditState` actor. + A progress spinner for `pip-audit`, using `rich.status`'s spinner support + under the hood. """ - def __init__(self, message: str = "", **kwargs: Dict[str, Any]): + def __init__(self, message: str = "") -> None: """ - Create a new `AuditSpinner`. - - `message` is the initial text that the progress spinner should display. - - Any remaining keyword arguments are forwarded onto the constructor of the underlying - `BaseSpinner` implementation. + Initialize the `AuditSpinner`. """ - super().__init__(message=message, **kwargs) + self._console = Console() + # NOTE: audits can be quite fast, so we need a pretty high refresh rate here. + self._spinner = self._console.status(message, spinner="line", refresh_per_second=30) # Keep the target set to `None` to ensure that the logs don't get written until the spinner # has finished writing output, regardless of the capacity argument @@ -125,67 +119,19 @@ def __init__(self, message: str = "", **kwargs: Dict[str, Any]): ) self.prev_handlers: List[logging.Handler] = [] - def _writeln_truncated(self, line: str) -> None: - """ - Wraps `BaseSpinner.writeln`, providing reasonable truncation behavior - when a line would otherwise overflow its terminal row and cause the progress - bar to break. - """ - if not (self.file and self.is_tty()): - return - - columns, _ = os.get_terminal_size(self.file.fileno()) - if columns > 4 and len(line) >= columns: - line = f"{line[0:columns - 4]} ..." - else: - line = line[0:columns] - - self.writeln(line) - - def update(self) -> None: - """ - Update the progress spinner. - - This method is overriden from `BaseSpinner` to customize the appearance of the spinner and - should not be called directly. - """ - i = self.index % len(self.phases) - line = f"{self.phases[i]} {self.message}" - self._writeln_truncated(line) - - def finish(self) -> None: - """ - Finish the progress spinner. - - This method is overridden from `BaseSpinner` to customize the spinner's termination - behavior: instead of finishing by printing a newline and leaving the last spinner state - on the terminal, we clear the spinner entirely and reset the line's state, leaving - no trace of the spinner at all. - """ - self.writeln("") - self.file.write("\r") - - # `BaseSpinner` normally re-reveals the cursor as part of `finish()` or - # `__del__`, but we override `finish()` and `__del__` isn't reliably - # invoked on context exit. So we do it manually here. - self.file.write(SHOW_CURSOR) - self.file.flush() - def update_state(self, message: str) -> None: """ - Update the state message for the progress spinner. - - This method is overriden from `AuditState` to update the spinner with feedback from the API - and should not be called directly. + Update the spinner's state. """ - self.message = message - self.next() + + self._spinner.update(message) def initialize(self) -> None: """ Redirect logging to an in-memory log handler so that it doesn't get mixed in with the spinner output. """ + # Remove all existing log handlers # # We're recording them here since we'll want to restore them once the spinner falls out of @@ -199,12 +145,15 @@ def initialize(self) -> None: # Redirect logging to our in-memory handler that will buffer the log lines root_logger.addHandler(self.log_handler) + self._spinner.start() + def finalize(self) -> None: """ Cleanup the spinner output so it doesn't get combined with subsequent `stderr` output and flush any logs that were recorded while the spinner was active. """ - self.finish() + + self._spinner.stop() # Now that the spinner is complete, flush the logs root_logger = logging.root diff --git a/pyproject.toml b/pyproject.toml index 8d2e65f0..e5c09eed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,8 @@ dependencies = [ "html5lib>=1.1", "packaging>=21.0.0", "pip-api>=0.0.28", - "progress>=1.6", "resolvelib>=0.8.0", + "rich>=12.4", ] requires-python = ">=3.7"