mirror of
https://github.com/Garmelon/PFERD.git
synced 2023-12-21 10:23:01 +01:00
112 lines
2.9 KiB
Python
112 lines
2.9 KiB
Python
"""
|
|
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)
|