Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

cli, state: replace progress spinner with rich.status #283

Merged
merged 2 commits into from
May 24, 2022
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pip_audit/_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
87 changes: 18 additions & 69 deletions pip_audit/_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

N.B.: If we wanted to be extra fancy, we could allow users to configure the spinner=... here, rather than using the same "line" spinner that progress uses.

Rich has a decent collection of spinners to choose from:

https://github.com/Textualize/rich/blob/aa7926c1431eebfb2ccaab9f3b63a4ac6cd8dfe6/rich/_spinners.py#L22


# 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
Expand All @@ -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
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down