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
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),
    transient=True
)


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

        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)