progress_bar.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224
  1. import math
  2. from functools import lru_cache
  3. from time import monotonic
  4. from typing import Iterable, List, Optional
  5. from .color import Color, blend_rgb
  6. from .color_triplet import ColorTriplet
  7. from .console import Console, ConsoleOptions, RenderResult
  8. from .jupyter import JupyterMixin
  9. from .measure import Measurement
  10. from .segment import Segment
  11. from .style import Style, StyleType
  12. # Number of characters before 'pulse' animation repeats
  13. PULSE_SIZE = 20
  14. class ProgressBar(JupyterMixin):
  15. """Renders a (progress) bar. Used by rich.progress.
  16. Args:
  17. total (float, optional): Number of steps in the bar. Defaults to 100. Set to None to render a pulsing animation.
  18. completed (float, optional): Number of steps completed. Defaults to 0.
  19. width (int, optional): Width of the bar, or ``None`` for maximum width. Defaults to None.
  20. pulse (bool, optional): Enable pulse effect. Defaults to False. Will pulse if a None total was passed.
  21. style (StyleType, optional): Style for the bar background. Defaults to "bar.back".
  22. complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".
  23. finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".
  24. pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".
  25. animation_time (Optional[float], optional): Time in seconds to use for animation, or None to use system time.
  26. """
  27. def __init__(
  28. self,
  29. total: Optional[float] = 100.0,
  30. completed: float = 0,
  31. width: Optional[int] = None,
  32. pulse: bool = False,
  33. style: StyleType = "bar.back",
  34. complete_style: StyleType = "bar.complete",
  35. finished_style: StyleType = "bar.finished",
  36. pulse_style: StyleType = "bar.pulse",
  37. animation_time: Optional[float] = None,
  38. ):
  39. self.total = total
  40. self.completed = completed
  41. self.width = width
  42. self.pulse = pulse
  43. self.style = style
  44. self.complete_style = complete_style
  45. self.finished_style = finished_style
  46. self.pulse_style = pulse_style
  47. self.animation_time = animation_time
  48. self._pulse_segments: Optional[List[Segment]] = None
  49. def __repr__(self) -> str:
  50. return f"<Bar {self.completed!r} of {self.total!r}>"
  51. @property
  52. def percentage_completed(self) -> Optional[float]:
  53. """Calculate percentage complete."""
  54. if self.total is None:
  55. return None
  56. completed = (self.completed / self.total) * 100.0
  57. completed = min(100, max(0.0, completed))
  58. return completed
  59. @lru_cache(maxsize=16)
  60. def _get_pulse_segments(
  61. self,
  62. fore_style: Style,
  63. back_style: Style,
  64. color_system: str,
  65. no_color: bool,
  66. ascii: bool = False,
  67. ) -> List[Segment]:
  68. """Get a list of segments to render a pulse animation.
  69. Returns:
  70. List[Segment]: A list of segments, one segment per character.
  71. """
  72. bar = "-" if ascii else "━"
  73. segments: List[Segment] = []
  74. if color_system not in ("standard", "eight_bit", "truecolor") or no_color:
  75. segments += [Segment(bar, fore_style)] * (PULSE_SIZE // 2)
  76. segments += [Segment(" " if no_color else bar, back_style)] * (
  77. PULSE_SIZE - (PULSE_SIZE // 2)
  78. )
  79. return segments
  80. append = segments.append
  81. fore_color = (
  82. fore_style.color.get_truecolor()
  83. if fore_style.color
  84. else ColorTriplet(255, 0, 255)
  85. )
  86. back_color = (
  87. back_style.color.get_truecolor()
  88. if back_style.color
  89. else ColorTriplet(0, 0, 0)
  90. )
  91. cos = math.cos
  92. pi = math.pi
  93. _Segment = Segment
  94. _Style = Style
  95. from_triplet = Color.from_triplet
  96. for index in range(PULSE_SIZE):
  97. position = index / PULSE_SIZE
  98. fade = 0.5 + cos((position * pi * 2)) / 2.0
  99. color = blend_rgb(fore_color, back_color, cross_fade=fade)
  100. append(_Segment(bar, _Style(color=from_triplet(color))))
  101. return segments
  102. def update(self, completed: float, total: Optional[float] = None) -> None:
  103. """Update progress with new values.
  104. Args:
  105. completed (float): Number of steps completed.
  106. total (float, optional): Total number of steps, or ``None`` to not change. Defaults to None.
  107. """
  108. self.completed = completed
  109. self.total = total if total is not None else self.total
  110. def _render_pulse(
  111. self, console: Console, width: int, ascii: bool = False
  112. ) -> Iterable[Segment]:
  113. """Renders the pulse animation.
  114. Args:
  115. console (Console): Console instance.
  116. width (int): Width in characters of pulse animation.
  117. Returns:
  118. RenderResult: [description]
  119. Yields:
  120. Iterator[Segment]: Segments to render pulse
  121. """
  122. fore_style = console.get_style(self.pulse_style, default="white")
  123. back_style = console.get_style(self.style, default="black")
  124. pulse_segments = self._get_pulse_segments(
  125. fore_style, back_style, console.color_system, console.no_color, ascii=ascii
  126. )
  127. segment_count = len(pulse_segments)
  128. current_time = (
  129. monotonic() if self.animation_time is None else self.animation_time
  130. )
  131. segments = pulse_segments * (int(width / segment_count) + 2)
  132. offset = int(-current_time * 15) % segment_count
  133. segments = segments[offset : offset + width]
  134. yield from segments
  135. def __rich_console__(
  136. self, console: Console, options: ConsoleOptions
  137. ) -> RenderResult:
  138. width = min(self.width or options.max_width, options.max_width)
  139. ascii = options.legacy_windows or options.ascii_only
  140. should_pulse = self.pulse or self.total is None
  141. if should_pulse:
  142. yield from self._render_pulse(console, width, ascii=ascii)
  143. return
  144. completed: Optional[float] = (
  145. min(self.total, max(0, self.completed)) if self.total is not None else None
  146. )
  147. bar = "-" if ascii else "━"
  148. half_bar_right = " " if ascii else "╸"
  149. half_bar_left = " " if ascii else "╺"
  150. complete_halves = (
  151. int(width * 2 * completed / self.total)
  152. if self.total and completed is not None
  153. else width * 2
  154. )
  155. bar_count = complete_halves // 2
  156. half_bar_count = complete_halves % 2
  157. style = console.get_style(self.style)
  158. is_finished = self.total is None or self.completed >= self.total
  159. complete_style = console.get_style(
  160. self.finished_style if is_finished else self.complete_style
  161. )
  162. _Segment = Segment
  163. if bar_count:
  164. yield _Segment(bar * bar_count, complete_style)
  165. if half_bar_count:
  166. yield _Segment(half_bar_right * half_bar_count, complete_style)
  167. if not console.no_color:
  168. remaining_bars = width - bar_count - half_bar_count
  169. if remaining_bars and console.color_system is not None:
  170. if not half_bar_count and bar_count:
  171. yield _Segment(half_bar_left, style)
  172. remaining_bars -= 1
  173. if remaining_bars:
  174. yield _Segment(bar * remaining_bars, style)
  175. def __rich_measure__(
  176. self, console: Console, options: ConsoleOptions
  177. ) -> Measurement:
  178. return (
  179. Measurement(self.width, self.width)
  180. if self.width is not None
  181. else Measurement(4, options.max_width)
  182. )
  183. if __name__ == "__main__": # pragma: no cover
  184. console = Console()
  185. bar = ProgressBar(width=50, total=100)
  186. import time
  187. console.show_cursor(False)
  188. for n in range(0, 101, 1):
  189. bar.update(n)
  190. console.print(bar)
  191. console.file.write("\r")
  192. time.sleep(0.05)
  193. console.show_cursor(True)
  194. console.print()