progress.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. from __future__ import annotations
  2. from typing import TYPE_CHECKING, Any, Dict, List, Optional
  3. from rich.console import Console, RenderableType
  4. from rich.live import Live
  5. from rich.text import Text
  6. from .element import Element
  7. if TYPE_CHECKING:
  8. from .styles.base import BaseStyle
  9. class ProgressLine(Element):
  10. def __init__(self, text: str | Text, parent: Progress):
  11. self.text = text
  12. self.parent = parent
  13. class Progress(Live, Element):
  14. current_message: str | Text
  15. def __init__(
  16. self,
  17. title: str,
  18. style: Optional[BaseStyle] = None,
  19. console: Optional[Console] = None,
  20. transient: bool = False,
  21. transient_on_error: bool = False,
  22. inline_logs: bool = False,
  23. lines_to_show: int = -1,
  24. **metadata: Dict[Any, Any],
  25. ) -> None:
  26. self.title = title
  27. self.current_message = title
  28. self.is_error = False
  29. self._transient_on_error = transient_on_error
  30. self._inline_logs = inline_logs
  31. self.lines_to_show = lines_to_show
  32. self.logs: List[ProgressLine] = []
  33. self.metadata = metadata
  34. self._cancelled = False
  35. Element.__init__(self, style=style)
  36. super().__init__(console=console, refresh_per_second=8, transient=transient)
  37. # TODO: remove this once rich uses "Self"
  38. def __enter__(self) -> "Progress":
  39. self.start(refresh=self._renderable is not None)
  40. return self
  41. def get_renderable(self) -> RenderableType:
  42. return self.style.render_element(self, done=not self._started)
  43. def log(self, text: str | Text) -> None:
  44. if self._inline_logs:
  45. self.logs.append(ProgressLine(text, self))
  46. else:
  47. self.current_message = text
  48. def set_error(self, text: str) -> None:
  49. self.current_message = text
  50. self.is_error = True
  51. self.transient = self._transient_on_error