pferd/PFERD/progress.py

"""
A small progress bar implementation.
"""
import sys
from dataclasses import dataclass
from types import TracebackType
from typing import Optional, Type

import requests
from rich.console import Console, ConsoleOptions, Control, RenderResult
from rich.live_render import LiveRender
from rich.progress import (BarColumn, DownloadColumn, Progress, TaskID,
                           TextColumn, TimeRemainingColumn,
                           TransferSpeedColumn)

_progress: Progress = Progress(
    TextColumn("[bold blue]{task.fields[name]}", justify="right"),
    BarColumn(bar_width=None),
    "[progress.percentage]{task.percentage:>3.1f}%",
    "•",
    DownloadColumn(),
    "•",
    TransferSpeedColumn(),
    "•",
    TimeRemainingColumn(),
    console=Console(file=sys.stdout)
)


def size_from_headers(response: requests.Response) -> Optional[int]:
    """
    Return the size of the download based on the response headers.

    Arguments:
        response {requests.Response} -- the response

    Returns:
        Optional[int] -- the size
    """
    if "Content-Length" in response.headers:
        return int(response.headers["Content-Length"])
    return None


@dataclass
class ProgressSettings:
    """
    Settings you can pass to customize the progress bar.
    """
    name: str
    max_size: int


def progress_for(settings: Optional[ProgressSettings]) -> 'ProgressContextManager':
    """
    Returns a context manager that displays progress

    Returns:
        ProgressContextManager -- the progress manager
    """
    return ProgressContextManager(settings)


class _OneLineUp(LiveRender):
    """
    Render a control code for moving one line upwards.
    """

    def __init__(self) -> None:
        super().__init__("not rendered")

    def __console__(self, console: Console, options: ConsoleOptions) -> RenderResult:
        yield Control(f"\r\x1b[1A")


class ProgressContextManager:
    """
    A context manager used for displaying progress.
    """

    def __init__(self, settings: Optional[ProgressSettings]):
        self._settings = settings
        self._task_id: Optional[TaskID] = None

    def __enter__(self) -> 'ProgressContextManager':
        """Context manager entry function."""
        if not self._settings:
            return self

        _progress.start()
        self._task_id = _progress.add_task(
            self._settings.name,
            total=self._settings.max_size,
            name=self._settings.name
        )
        return self

    # pylint: disable=useless-return
    def __exit__(
            self,
            exc_type: Optional[Type[BaseException]],
            exc_value: Optional[BaseException],
            traceback: Optional[TracebackType],
    ) -> Optional[bool]:
        """Context manager exit function. Removes the task."""
        if self._task_id is None:
            return None

        _progress.remove_task(self._task_id)

        if len(_progress.task_ids) == 0:
            # We need to clean up after ourselves, as we were the last one
            _progress.stop()
            _progress.refresh()

            # And we existed, so remove the line above (remove_task leaves one behind)
            Console().print(_OneLineUp())

        return None

    def advance(self, amount: float) -> None:
        """
        Advances the progress bar.
        """
        if self._task_id is not None:
            _progress.advance(self._task_id, amount)