updates

Backend/venv/lib/python3.12/site-packages/pip_audit/_state.py

@@ -0,0 +1,274 @@
+"""
+Interfaces for for propagating feedback from the API to provide responsive progress indicators as
+well as a progress spinner implementation for use with CLI applications.
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from logging.handlers import MemoryHandler
+from typing import Any
+
+from rich.align import StyleType
+from rich.console import Console, Group, RenderableType
+from rich.live import Live
+from rich.panel import Panel
+from rich.status import Spinner
+
+
+class AuditState:
+    """
+    An object that handles abstract "updates" to `pip-audit`'s state.
+
+    Non-UI consumers of `pip-audit` (via `pip_audit`) should have no need for
+    this class, and can leave it as a default construction in whatever signatures
+    it appears in. Its primary use is internal and UI-specific: it exists solely
+    to give the CLI enough state for a responsive progress indicator during
+    user requests.
+    """
+
+    def __init__(self, *, members: Sequence[_StateActor] = []):
+        """
+        Create a new `AuditState` with the given member list.
+        """
+
+        self._members = members
+
+    def update_state(self, message: str, logs: str | None = None) -> None:
+        """
+        Called whenever `pip_audit`'s internal state changes in a way that's meaningful to
+        expose to a user.
+
+        `message` is the message to present to the user.
+        """
+
+        for member in self._members:
+            member.update_state(message, logs)
+
+    def initialize(self) -> None:
+        """
+        Called when `pip-audit`'s state is initializing.
+        """
+
+        for member in self._members:
+            member.initialize()
+
+    def finalize(self) -> None:
+        """
+        Called when `pip_audit`'s state is "done" changing.
+        """
+        for member in self._members:
+            member.finalize()
+
+    def __enter__(self) -> AuditState:  # pragma: no cover
+        """
+        Create an instance of the `pip-audit` state for usage within a `with` statement.
+        """
+
+        self.initialize()
+        return self
+
+    def __exit__(
+        self, _exc_type: Any, _exc_value: Any, _exc_traceback: Any
+    ) -> None:  # pragma: no cover
+        """
+        Helper to ensure `finalize` gets called when the `pip-audit` state falls out of scope of a
+        `with` statement.
+        """
+        self.finalize()
+
+
+class _StateActor(ABC):
+    @abstractmethod
+    def update_state(self, message: str, logs: str | None = None) -> None:
+        raise NotImplementedError  # pragma: no cover
+
+    @abstractmethod
+    def initialize(self) -> None:
+        """
+        Called when `pip-audit`'s state is initializing. Implementors should
+        override this to do nothing if their state management requires no
+        initialization step.
+        """
+        raise NotImplementedError  # pragma: no cover
+
+    @abstractmethod
+    def finalize(self) -> None:
+        """
+        Called when the overlaying `AuditState` is "done," i.e. `pip-audit`'s
+        state is done changing. Implementors should override this to do nothing
+        if their state management requires no finalization step.
+        """
+        raise NotImplementedError  # pragma: no cover
+
+
+class StatusLog:  # pragma: no cover
+    """
+    Displays a status indicator with an optional log panel to display logs
+    for external processes.
+
+    This code is based off of Rich's `Status` component:
+        https://github.com/Textualize/rich/blob/master/rich/status.py
+    """
+
+    # NOTE(alex): We limit the panel to 10 characters high and display the last 10 log lines.
+    # However, the panel won't display all 10 of those lines if some of the lines are long enough
+    # to wrap in the panel.
+    LOG_PANEL_HEIGHT = 10
+
+    def __init__(
+        self,
+        status: str,
+        *,
+        console: Console | None = None,
+        spinner: str = "dots",
+        spinner_style: StyleType = "status.spinner",
+        speed: float = 1.0,
+        refresh_per_second: float = 12.5,
+    ):
+        """
+        Construct a new `StatusLog`.
+
+        `status` is the status message to display next to the spinner.
+        `console` is the Rich console to display the log status in.
+        `spinner` is the name of the spinner animation (see python -m rich.spinner). Defaults to `dots`.
+        `spinner_style` is the style of the spinner. Defaults to `status.spinner`.
+        `speed` is the speed factor for the spinner animation. Defaults to 1.0.
+        `refresh_per_second` is the number of refreshes per second. Defaults to 12.5.
+        """
+
+        self._spinner = Spinner(spinner, text=status, style=spinner_style, speed=speed)
+        self._log_panel = Panel("", height=self.LOG_PANEL_HEIGHT)
+        self._live = Live(
+            self.renderable,
+            console=console,
+            refresh_per_second=refresh_per_second,
+            transient=True,
+        )
+
+    @property
+    def renderable(self) -> RenderableType:
+        """
+        Create a Rich renderable type for the log panel.
+
+        If the log panel contains text, we should create a group and place the
+        log panel underneath the spinner.
+        """
+
+        if self._log_panel.renderable:
+            return Group(self._spinner, self._log_panel)
+        return self._spinner
+
+    def update(
+        self,
+        status: str,
+        logs: str | None,
+    ) -> None:
+        """
+        Update status and logs.
+        """
+
+        if logs is None:
+            logs = ""
+        else:
+            # Limit the logging output to the 10 most recent lines.
+            logs = "\n".join(logs.splitlines()[-self.LOG_PANEL_HEIGHT :])
+        self._spinner.update(text=status)
+        self._log_panel.renderable = logs
+        self._live.update(self.renderable, refresh=True)
+
+    def start(self) -> None:
+        """
+        Start the status animation.
+        """
+
+        self._live.start()
+
+    def stop(self) -> None:
+        """
+        Stop the spinner animation.
+        """
+
+        self._live.stop()
+
+    def __rich__(self) -> RenderableType:
+        """
+        Convert to a Rich renderable type.
+        """
+
+        return self.renderable
+
+
+class AuditSpinner(_StateActor):  # pragma: no cover
+    """
+    A progress spinner for `pip-audit`, using `rich.status`'s spinner support
+    under the hood.
+    """
+
+    def __init__(self, message: str = "") -> None:
+        """
+        Initialize the `AuditSpinner`.
+        """
+
+        self._console = Console()
+        # NOTE: audits can be quite fast, so we need a pretty high refresh rate here.
+        self._spinner = StatusLog(
+            message, console=self._console, 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
+        self.log_handler = MemoryHandler(
+            0, flushLevel=logging.ERROR, target=None, flushOnClose=False
+        )
+        self.prev_handlers: list[logging.Handler] = []
+
+    def update_state(self, message: str, logs: str | None = None) -> None:
+        """
+        Update the spinner's state.
+        """
+
+        self._spinner.update(message, logs)
+
+    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
+        # scope
+        root_logger = logging.root
+        for handler in root_logger.handlers:
+            self.prev_handlers.append(handler)
+        for handler in self.prev_handlers:
+            root_logger.removeHandler(handler)
+
+        # 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._spinner.stop()
+
+        # Now that the spinner is complete, flush the logs
+        root_logger = logging.root
+        stream_handler = logging.StreamHandler()
+        stream_handler.setFormatter(logging.Formatter(logging.BASIC_FORMAT))
+        self.log_handler.setTarget(stream_handler)
+        self.log_handler.flush()
+
+        # Restore the original log handlers
+        root_logger.removeHandler(self.log_handler)
+        for handler in self.prev_handlers:
+            root_logger.addHandler(handler)