Add beautiful progressbar (#139)

* Add beautiful progressbar

* minor refinement

mim/utils/__init__.py

@@ -9,6 +9,7 @@
     USER,
     WHEEL_URL,
 )
+from .progress_bars import rich_progress_bar
 from .utils import (
     args2string,
     call_command,
@@ -86,4 +87,5 @@
     'get_package_info_from_pypi',
     'parse_home_page',
     'ensure_installation',
+    'rich_progress_bar',
 ]

mim/utils/progress_bars.py

@@ -0,0 +1,37 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# Copyright (c) https://github.com/pypa/pip
+# Modified from https://github.com/pypa/pip/blob/main/src/pip/_internal/cli/progress_bars.py  # noqa: E501
+
+from typing import Callable, Generator, Iterable, Iterator
+
+from rich.progress import (
+    BarColumn,
+    DownloadColumn,
+    Progress,
+    TextColumn,
+    TimeRemainingColumn,
+    TransferSpeedColumn,
+)
+
+DownloadProgressRenderer = Callable[[Iterable[bytes]], Iterator[bytes]]
+
+
+def rich_progress_bar(
+    iterable: Iterable[bytes],
+    size: int,
+) -> Generator[bytes, None, None]:
+    columns = (
+        TextColumn('[progress.description]{task.description}'),
+        BarColumn(),
+        DownloadColumn(binary_units=1024),
+        TransferSpeedColumn(),
+        TextColumn('eta'),
+        TimeRemainingColumn(),
+    )
+
+    progress = Progress(*columns, refresh_per_second=30)
+    task_id = progress.add_task('downloading', total=size)
+    with progress:
+        for chunk in iterable:
+            yield chunk
+            progress.update(task_id, advance=len(chunk))

mim/utils/utils.py

@@ -21,6 +21,7 @@
 from requests.models import Response
 
 from .default import PKG2PROJECT
+from .progress_bars import rich_progress_bar
 
 
 def parse_url(url: str) -> Tuple[str, str]:
@@ -172,13 +173,12 @@ def download_from_file(url: str,
     size = int(response.headers.get('content-length'))
     with open(dest_path, 'wb') as fw:
         content_iter = response.iter_content(chunk_size=1024)
-        with click.progressbar(content_iter, length=size / 1024) as chunks:
-            for chunk in chunks:
-                if chunk:
-                    fw.write(chunk)
-                    fw.flush()
-                    if hash_prefix is not None:
-                        sha256.update(chunk)
+        for chunk in rich_progress_bar(content_iter, size=size):
+            if chunk:
+                fw.write(chunk)
+                fw.flush()
+                if hash_prefix is not None:
+                    sha256.update(chunk)
 
     if hash_prefix is not None:
         digest = sha256.hexdigest()

requirements/install.txt

@@ -4,4 +4,5 @@ model-index
 pandas
 pip>=19.3
 requests
+rich
 tabulate