_state.py

"""
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)