57 lines
1.5 KiB
Python
57 lines
1.5 KiB
Python
from typing import Optional
|
|
|
|
from rich.console import Console
|
|
from rich.live import Live
|
|
from typing_extensions import Any, Literal
|
|
|
|
from .styles.base import BaseStyle
|
|
|
|
|
|
class Progress(Live):
|
|
def __init__(
|
|
self,
|
|
title: str,
|
|
style: Optional[BaseStyle] = None,
|
|
console: Optional[Console] = None,
|
|
transient: bool = False,
|
|
transient_on_error: bool = False,
|
|
) -> None:
|
|
self.current_message = title
|
|
self.style = style
|
|
self.is_error = False
|
|
self._transient_on_error = transient_on_error
|
|
|
|
super().__init__(console=console, refresh_per_second=8, transient=transient)
|
|
|
|
# TODO: remove this once rich uses "Self"
|
|
def __enter__(self) -> "Progress":
|
|
self.start(refresh=self._renderable is not None)
|
|
|
|
return self
|
|
|
|
def get_renderable(self) -> Any:
|
|
current_message = self.current_message
|
|
|
|
if not self.style:
|
|
return current_message
|
|
|
|
animation_status: Literal["started", "stopped", "error"] = (
|
|
"started" if self._started else "stopped"
|
|
)
|
|
|
|
if self.is_error:
|
|
animation_status = "error"
|
|
|
|
return self.style.with_decoration(
|
|
current_message,
|
|
animation_status=animation_status,
|
|
)
|
|
|
|
def log(self, text: str) -> None:
|
|
self.current_message = text
|
|
|
|
def set_error(self, text: str) -> None:
|
|
self.current_message = text
|
|
self.is_error = True
|
|
self.transient = self._transient_on_error
|