progress_bars.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. import functools
  2. from typing import Callable, Generator, Iterable, Iterator, Optional, Tuple
  3. from pip._vendor.rich.progress import (
  4. BarColumn,
  5. DownloadColumn,
  6. FileSizeColumn,
  7. Progress,
  8. ProgressColumn,
  9. SpinnerColumn,
  10. TextColumn,
  11. TimeElapsedColumn,
  12. TimeRemainingColumn,
  13. TransferSpeedColumn,
  14. )
  15. from pip._internal.utils.logging import get_indentation
  16. DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]]
  17. def _rich_progress_bar(
  18. iterable: Iterable[bytes],
  19. *,
  20. bar_type: str,
  21. size: int,
  22. ) -> Generator[bytes, None, None]:
  23. assert bar_type == "on", "This should only be used in the default mode."
  24. if not size:
  25. total = float("inf")
  26. columns: Tuple[ProgressColumn, ...] = (
  27. TextColumn("[progress.description]{task.description}"),
  28. SpinnerColumn("line", speed=1.5),
  29. FileSizeColumn(),
  30. TransferSpeedColumn(),
  31. TimeElapsedColumn(),
  32. )
  33. else:
  34. total = size
  35. columns = (
  36. TextColumn("[progress.description]{task.description}"),
  37. BarColumn(),
  38. DownloadColumn(),
  39. TransferSpeedColumn(),
  40. TextColumn("eta"),
  41. TimeRemainingColumn(),
  42. )
  43. progress = Progress(*columns, refresh_per_second=30)
  44. task_id = progress.add_task(" " * (get_indentation() + 2), total=total)
  45. with progress:
  46. for chunk in iterable:
  47. yield chunk
  48. progress.update(task_id, advance=len(chunk))
  49. def get_download_progress_renderer(
  50. *, bar_type: str, size: Optional[int] = None
  51. ) -> DownloadProgressRenderer:
  52. """Get an object that can be used to render the download progress.
  53. Returns a callable, that takes an iterable to "wrap".
  54. """
  55. if bar_type == "on":
  56. return functools.partial(_rich_progress_bar, bar_type=bar_type, size=size)
  57. else:
  58. return iter # no-op, when passed an iterator