progress_bars.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321
  1. import functools
  2. import itertools
  3. import sys
  4. from signal import SIGINT, default_int_handler, signal
  5. from typing import Any, Callable, Iterator, Optional, Tuple
  6. from pip._vendor.progress.bar import Bar, FillingCirclesBar, IncrementalBar
  7. from pip._vendor.progress.spinner import Spinner
  8. from pip._vendor.rich.progress import (
  9. BarColumn,
  10. DownloadColumn,
  11. FileSizeColumn,
  12. Progress,
  13. ProgressColumn,
  14. SpinnerColumn,
  15. TextColumn,
  16. TimeElapsedColumn,
  17. TimeRemainingColumn,
  18. TransferSpeedColumn,
  19. )
  20. from pip._internal.utils.compat import WINDOWS
  21. from pip._internal.utils.logging import get_indentation
  22. from pip._internal.utils.misc import format_size
  23. try:
  24. from pip._vendor import colorama
  25. # Lots of different errors can come from this, including SystemError and
  26. # ImportError.
  27. except Exception:
  28. colorama = None
  29. DownloadProgressRenderer = Callable[[Iterator[bytes]], Iterator[bytes]]
  30. def _select_progress_class(preferred: Bar, fallback: Bar) -> Bar:
  31. encoding = getattr(preferred.file, "encoding", None)
  32. # If we don't know what encoding this file is in, then we'll just assume
  33. # that it doesn't support unicode and use the ASCII bar.
  34. if not encoding:
  35. return fallback
  36. # Collect all of the possible characters we want to use with the preferred
  37. # bar.
  38. characters = [
  39. getattr(preferred, "empty_fill", ""),
  40. getattr(preferred, "fill", ""),
  41. ]
  42. characters += list(getattr(preferred, "phases", []))
  43. # Try to decode the characters we're using for the bar using the encoding
  44. # of the given file, if this works then we'll assume that we can use the
  45. # fancier bar and if not we'll fall back to the plaintext bar.
  46. try:
  47. "".join(characters).encode(encoding)
  48. except UnicodeEncodeError:
  49. return fallback
  50. else:
  51. return preferred
  52. _BaseBar: Any = _select_progress_class(IncrementalBar, Bar)
  53. class InterruptibleMixin:
  54. """
  55. Helper to ensure that self.finish() gets called on keyboard interrupt.
  56. This allows downloads to be interrupted without leaving temporary state
  57. (like hidden cursors) behind.
  58. This class is similar to the progress library's existing SigIntMixin
  59. helper, but as of version 1.2, that helper has the following problems:
  60. 1. It calls sys.exit().
  61. 2. It discards the existing SIGINT handler completely.
  62. 3. It leaves its own handler in place even after an uninterrupted finish,
  63. which will have unexpected delayed effects if the user triggers an
  64. unrelated keyboard interrupt some time after a progress-displaying
  65. download has already completed, for example.
  66. """
  67. def __init__(self, *args: Any, **kwargs: Any) -> None:
  68. """
  69. Save the original SIGINT handler for later.
  70. """
  71. # https://github.com/python/mypy/issues/5887
  72. super().__init__(*args, **kwargs) # type: ignore
  73. self.original_handler = signal(SIGINT, self.handle_sigint)
  74. # If signal() returns None, the previous handler was not installed from
  75. # Python, and we cannot restore it. This probably should not happen,
  76. # but if it does, we must restore something sensible instead, at least.
  77. # The least bad option should be Python's default SIGINT handler, which
  78. # just raises KeyboardInterrupt.
  79. if self.original_handler is None:
  80. self.original_handler = default_int_handler
  81. def finish(self) -> None:
  82. """
  83. Restore the original SIGINT handler after finishing.
  84. This should happen regardless of whether the progress display finishes
  85. normally, or gets interrupted.
  86. """
  87. super().finish() # type: ignore
  88. signal(SIGINT, self.original_handler)
  89. def handle_sigint(self, signum, frame): # type: ignore
  90. """
  91. Call self.finish() before delegating to the original SIGINT handler.
  92. This handler should only be in place while the progress display is
  93. active.
  94. """
  95. self.finish()
  96. self.original_handler(signum, frame)
  97. class SilentBar(Bar):
  98. def update(self) -> None:
  99. pass
  100. class BlueEmojiBar(IncrementalBar):
  101. suffix = "%(percent)d%%"
  102. bar_prefix = " "
  103. bar_suffix = " "
  104. phases = ("\U0001F539", "\U0001F537", "\U0001F535")
  105. class DownloadProgressMixin:
  106. def __init__(self, *args: Any, **kwargs: Any) -> None:
  107. # https://github.com/python/mypy/issues/5887
  108. super().__init__(*args, **kwargs) # type: ignore
  109. self.message: str = (" " * (get_indentation() + 2)) + self.message
  110. @property
  111. def downloaded(self) -> str:
  112. return format_size(self.index) # type: ignore
  113. @property
  114. def download_speed(self) -> str:
  115. # Avoid zero division errors...
  116. if self.avg == 0.0: # type: ignore
  117. return "..."
  118. return format_size(1 / self.avg) + "/s" # type: ignore
  119. @property
  120. def pretty_eta(self) -> str:
  121. if self.eta: # type: ignore
  122. return f"eta {self.eta_td}" # type: ignore
  123. return ""
  124. def iter(self, it): # type: ignore
  125. for x in it:
  126. yield x
  127. # B305 is incorrectly raised here
  128. # https://github.com/PyCQA/flake8-bugbear/issues/59
  129. self.next(len(x)) # noqa: B305
  130. self.finish()
  131. class WindowsMixin:
  132. def __init__(self, *args: Any, **kwargs: Any) -> None:
  133. # The Windows terminal does not support the hide/show cursor ANSI codes
  134. # even with colorama. So we'll ensure that hide_cursor is False on
  135. # Windows.
  136. # This call needs to go before the super() call, so that hide_cursor
  137. # is set in time. The base progress bar class writes the "hide cursor"
  138. # code to the terminal in its init, so if we don't set this soon
  139. # enough, we get a "hide" with no corresponding "show"...
  140. if WINDOWS and self.hide_cursor: # type: ignore
  141. self.hide_cursor = False
  142. # https://github.com/python/mypy/issues/5887
  143. super().__init__(*args, **kwargs) # type: ignore
  144. # Check if we are running on Windows and we have the colorama module,
  145. # if we do then wrap our file with it.
  146. if WINDOWS and colorama:
  147. self.file = colorama.AnsiToWin32(self.file) # type: ignore
  148. # The progress code expects to be able to call self.file.isatty()
  149. # but the colorama.AnsiToWin32() object doesn't have that, so we'll
  150. # add it.
  151. self.file.isatty = lambda: self.file.wrapped.isatty()
  152. # The progress code expects to be able to call self.file.flush()
  153. # but the colorama.AnsiToWin32() object doesn't have that, so we'll
  154. # add it.
  155. self.file.flush = lambda: self.file.wrapped.flush()
  156. class BaseDownloadProgressBar(WindowsMixin, InterruptibleMixin, DownloadProgressMixin):
  157. file = sys.stdout
  158. message = "%(percent)d%%"
  159. suffix = "%(downloaded)s %(download_speed)s %(pretty_eta)s"
  160. class DefaultDownloadProgressBar(BaseDownloadProgressBar, _BaseBar):
  161. pass
  162. class DownloadSilentBar(BaseDownloadProgressBar, SilentBar):
  163. pass
  164. class DownloadBar(BaseDownloadProgressBar, Bar):
  165. pass
  166. class DownloadFillingCirclesBar(BaseDownloadProgressBar, FillingCirclesBar):
  167. pass
  168. class DownloadBlueEmojiProgressBar(BaseDownloadProgressBar, BlueEmojiBar):
  169. pass
  170. class DownloadProgressSpinner(
  171. WindowsMixin, InterruptibleMixin, DownloadProgressMixin, Spinner
  172. ):
  173. file = sys.stdout
  174. suffix = "%(downloaded)s %(download_speed)s"
  175. def next_phase(self) -> str:
  176. if not hasattr(self, "_phaser"):
  177. self._phaser = itertools.cycle(self.phases)
  178. return next(self._phaser)
  179. def update(self) -> None:
  180. message = self.message % self
  181. phase = self.next_phase()
  182. suffix = self.suffix % self
  183. line = "".join(
  184. [
  185. message,
  186. " " if message else "",
  187. phase,
  188. " " if suffix else "",
  189. suffix,
  190. ]
  191. )
  192. self.writeln(line)
  193. BAR_TYPES = {
  194. "off": (DownloadSilentBar, DownloadSilentBar),
  195. "on": (DefaultDownloadProgressBar, DownloadProgressSpinner),
  196. "ascii": (DownloadBar, DownloadProgressSpinner),
  197. "pretty": (DownloadFillingCirclesBar, DownloadProgressSpinner),
  198. "emoji": (DownloadBlueEmojiProgressBar, DownloadProgressSpinner),
  199. }
  200. def _legacy_progress_bar(
  201. progress_bar: str, max: Optional[int]
  202. ) -> DownloadProgressRenderer:
  203. if max is None or max == 0:
  204. return BAR_TYPES[progress_bar][1]().iter # type: ignore
  205. else:
  206. return BAR_TYPES[progress_bar][0](max=max).iter
  207. #
  208. # Modern replacement, for our legacy progress bars.
  209. #
  210. def _rich_progress_bar(
  211. iterable: Iterator[bytes],
  212. *,
  213. bar_type: str,
  214. size: int,
  215. ) -> Iterator[bytes]:
  216. assert bar_type == "on", "This should only be used in the default mode."
  217. if not size:
  218. total = float("inf")
  219. columns: Tuple[ProgressColumn, ...] = (
  220. TextColumn("[progress.description]{task.description}"),
  221. SpinnerColumn("line", speed=1.5),
  222. FileSizeColumn(),
  223. TransferSpeedColumn(),
  224. TimeElapsedColumn(),
  225. )
  226. else:
  227. total = size
  228. columns = (
  229. TextColumn("[progress.description]{task.description}"),
  230. BarColumn(),
  231. DownloadColumn(),
  232. TransferSpeedColumn(),
  233. TextColumn("eta"),
  234. TimeRemainingColumn(),
  235. )
  236. progress = Progress(*columns, refresh_per_second=30)
  237. task_id = progress.add_task(" " * (get_indentation() + 2), total=total)
  238. with progress:
  239. for chunk in iterable:
  240. yield chunk
  241. progress.update(task_id, advance=len(chunk))
  242. def get_download_progress_renderer(
  243. *, bar_type: str, size: Optional[int] = None
  244. ) -> DownloadProgressRenderer:
  245. """Get an object that can be used to render the download progress.
  246. Returns a callable, that takes an iterable to "wrap".
  247. """
  248. if bar_type == "on":
  249. return functools.partial(_rich_progress_bar, bar_type=bar_type, size=size)
  250. elif bar_type == "off":
  251. return iter # no-op, when passed an iterator
  252. else:
  253. return _legacy_progress_bar(bar_type, size)