From 90d97cd2f70fdc6cc3e50db981647ccf29328194 Mon Sep 17 00:00:00 2001
From: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Date: Fri, 24 Jun 2022 14:01:17 +0800
Subject: [PATCH] Add beautiful progressbar (#139)

* Add beautiful progressbar

* minor refinement
---
 mim/utils/__init__.py      |  2 ++
 mim/utils/progress_bars.py | 37 +++++++++++++++++++++++++++++++++++++
 mim/utils/utils.py         | 14 +++++++-------
 requirements/install.txt   |  1 +
 4 files changed, 47 insertions(+), 7 deletions(-)
 create mode 100644 mim/utils/progress_bars.py

diff --git a/mim/utils/__init__.py b/mim/utils/__init__.py
index 98cd7f9..a45f2af 100644
--- a/mim/utils/__init__.py
+++ b/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',
 ]
diff --git a/mim/utils/progress_bars.py b/mim/utils/progress_bars.py
new file mode 100644
index 0000000..3505533
--- /dev/null
+++ b/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))
diff --git a/mim/utils/utils.py b/mim/utils/utils.py
index 485e021..8877b47 100644
--- a/mim/utils/utils.py
+++ b/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()
diff --git a/requirements/install.txt b/requirements/install.txt
index 74d075c..147ca51 100644
--- a/requirements/install.txt
+++ b/requirements/install.txt
@@ -4,4 +4,5 @@ model-index
 pandas
 pip>=19.3
 requests
+rich
 tabulate