progress_bars.py 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. import functools
  2. import sys
  3. from typing import Callable, Generator, Iterable, Iterator, Optional, Tuple
  4. from pip._vendor.rich.progress import (
  5. BarColumn,
  6. DownloadColumn,
  7. FileSizeColumn,
  8. Progress,
  9. ProgressColumn,
  10. SpinnerColumn,
  11. TextColumn,
  12. TimeElapsedColumn,
  13. TimeRemainingColumn,
  14. TransferSpeedColumn,
  15. )
  16. from pip._internal.cli.spinners import RateLimiter
  17. from pip._internal.utils.logging import get_indentation
  18. DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]]
  19. def _rich_progress_bar(
  20. iterable: Iterable[bytes],
  21. *,
  22. bar_type: str,
  23. size: int,
  24. ) -> Generator[bytes, None, None]:
  25. assert bar_type == "on", "This should only be used in the default mode."
  26. if not size:
  27. total = float("inf")
  28. columns: Tuple[ProgressColumn, ...] = (
  29. TextColumn("[progress.description]{task.description}"),
  30. SpinnerColumn("line", speed=1.5),
  31. FileSizeColumn(),
  32. TransferSpeedColumn(),
  33. TimeElapsedColumn(),
  34. )
  35. else:
  36. total = size
  37. columns = (
  38. TextColumn("[progress.description]{task.description}"),
  39. BarColumn(),
  40. DownloadColumn(),
  41. TransferSpeedColumn(),
  42. TextColumn("eta"),
  43. TimeRemainingColumn(),
  44. )
  45. progress = Progress(*columns, refresh_per_second=5)
  46. task_id = progress.add_task(" " * (get_indentation() + 2), total=total)
  47. with progress:
  48. for chunk in iterable:
  49. yield chunk
  50. progress.update(task_id, advance=len(chunk))
  51. def _raw_progress_bar(
  52. iterable: Iterable[bytes],
  53. *,
  54. size: Optional[int],
  55. ) -> Generator[bytes, None, None]:
  56. def write_progress(current: int, total: int) -> None:
  57. sys.stdout.write("Progress %d of %d\n" % (current, total))
  58. sys.stdout.flush()
  59. current = 0
  60. total = size or 0
  61. rate_limiter = RateLimiter(0.25)
  62. write_progress(current, total)
  63. for chunk in iterable:
  64. current += len(chunk)
  65. if rate_limiter.ready() or current == total:
  66. write_progress(current, total)
  67. rate_limiter.reset()
  68. yield chunk
  69. def get_download_progress_renderer(
  70. *, bar_type: str, size: Optional[int] = None
  71. ) -> DownloadProgressRenderer:
  72. """Get an object that can be used to render the download progress.
  73. Returns a callable, that takes an iterable to "wrap".
  74. """
  75. if bar_type == "on":
  76. return functools.partial(_rich_progress_bar, bar_type=bar_type, size=size)
  77. elif bar_type == "raw":
  78. return functools.partial(_raw_progress_bar, size=size)
  79. else:
  80. return iter # no-op, when passed an iterator