httpie-cli/httpie/output/ui/rich_progress.py

from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

from httpie.context import Environment

if TYPE_CHECKING:
    from rich.console import Console


@dataclass
class BaseDisplay:
    env: Environment

    def start(
        self, *, total: Optional[float], at: float, description: str
    ) -> None:
        ...

    def update(self, steps: float) -> None:
        ...

    def stop(self, time_spent: float) -> None:
        ...

    @property
    def console(self) -> 'Console':
        """Returns the default console to be used with displays (stderr)."""
        return self.env.rich_error_console

    def _print_summary(
        self, is_finished: bool, observed_steps: int, time_spent: float
    ):
        from rich import filesize

        if is_finished:
            verb = 'Done'
        else:
            verb = 'Interrupted'

        total_size = filesize.decimal(observed_steps)
        avg_speed = filesize.decimal(observed_steps / time_spent)

        minutes, seconds = divmod(time_spent, 60)
        hours, minutes = divmod(int(minutes), 60)
        if hours:
            total_time = f'{hours:d}:{minutes:02d}:{seconds:0.5f}'
        else:
            total_time = f'{minutes:02d}:{seconds:0.5f}'

        self.console.print(
            f'[progress.description]{verb}. {total_size} in {total_time} ({avg_speed}/s)'
        )


class DummyDisplay(BaseDisplay):
    """
    A dummy display object to be used when the progress bars,
    spinners etc. are disabled globally (or during tests).
    """


class StatusDisplay(BaseDisplay):
    def start(
        self, *, total: Optional[float], at: float, description: str
    ) -> None:
        self.observed = at
        self.description = (
            f'[progress.description]{description}[/progress.description]'
        )

        self.status = self.console.status(self.description, spinner='line')
        self.status.start()

    def update(self, steps: float) -> None:
        from rich import filesize

        self.observed += steps

        observed_amount, observed_unit = filesize.decimal(
            self.observed
        ).split()
        self.status.update(
            status=f'{self.description} [progress.download]{observed_amount}/? {observed_unit}[/progress.download]'
        )

    def stop(self, time_spent: float) -> None:
        self.status.stop()
        self.console.print(self.description)
        if time_spent:
            self._print_summary(
                is_finished=True,
                observed_steps=self.observed,
                time_spent=time_spent,
            )


class ProgressDisplay(BaseDisplay):
    def start(
        self, *, total: Optional[float], at: float, description: str
    ) -> None:
        from rich.progress import (
            Progress,
            BarColumn,
            DownloadColumn,
            TimeRemainingColumn,
            TransferSpeedColumn,
        )

        assert total is not None
        self.console.print(f'[progress.description]{description}')
        self.progress_bar = Progress(
            '[',
            BarColumn(),
            ']',
            '[progress.percentage]{task.percentage:>3.0f}%',
            '(',
            DownloadColumn(),
            ')',
            TimeRemainingColumn(),
            TransferSpeedColumn(),
            console=self.console,
            transient=True,
        )
        self.progress_bar.start()
        self.transfer_task = self.progress_bar.add_task(
            description, completed=at, total=total
        )

    def update(self, steps: float) -> None:
        self.progress_bar.advance(self.transfer_task, steps)

    def stop(self, time_spent: Optional[float]) -> None:
        self.progress_bar.stop()

        if time_spent:
            [task] = self.progress_bar.tasks
            self._print_summary(
                is_finished=task.finished,
                observed_steps=task.completed,
                time_spent=time_spent,
            )