Migrate progress bars from tqdm to rich.progress

Author: Jackenmen

redbot/cogs/downloader/repo_manager.py

@@ -30,7 +30,7 @@
 
 import discord
 from redbot.core import data_manager, commands, Config
-from redbot.core.utils._internal_utils import safe_delete
+from redbot.core.utils._internal_utils import detailed_progress, safe_delete
 from redbot.core.i18n import Translator
 
 from . import errors
@@ -1267,28 +1267,27 @@ async def _restore_from_backup(self) -> None:
         with open(data_manager.cog_data_path(self) / "repos.json") as fp:
             raw_repos = json.load(fp)
 
-        from tqdm import tqdm
-
-        progress_bar = tqdm(raw_repos, desc="Downloading repos", unit="repo", dynamic_ncols=True)
-
-        for repo_data in progress_bar:
-            repo_url = repo_data["url"]
-            repo_name = repo_data["name"]
-            repo_branch = repo_data["branch"]
-            try:
-                await self.add_repo(repo_url, repo_name, repo_branch)
-            except errors.CloningError:
-                log.exception(
-                    "Something went wrong whilst cloning %s (to branch: %s)",
-                    repo_url,
-                    repo_branch,
-                )
-            except OSError:
-                log.exception(
-                    "Something went wrong trying to add repo %s under name %s",
-                    repo_url,
-                    repo_name,
-                )
+        with detailed_progress(unit="repos") as progress:
+            task_id = progress.add_task("Adding repos", total=len(raw_repos))
+            for idx, repo_data in enumerate(raw_repos):
+                repo_url = repo_data["url"]
+                repo_name = repo_data["name"]
+                repo_branch = repo_data["branch"]
+                progress.update(task_id, completed=idx, description=f"Adding {repo_name!r} repo")
+                try:
+                    await self.add_repo(repo_url, repo_name, repo_branch)
+                except errors.CloningError:
+                    log.exception(
+                        "Something went wrong whilst cloning %s (to branch: %s)",
+                        repo_url,
+                        repo_branch,
+                    )
+                except OSError:
+                    log.exception(
+                        "Something went wrong trying to add repo %s under name %s",
+                        repo_url,
+                        repo_name,
+                    )
 
         from .downloader import Downloader
 

redbot/core/utils/_internal_utils.py

@@ -36,10 +36,9 @@
 import discord
 from packaging.requirements import Requirement
 import rapidfuzz
-from rich.progress import ProgressColumn
-from rich.progress_bar import ProgressBar
+import rich.progress
+from rich.text import Text
 from red_commons.logging import VERBOSE, TRACE
-from tqdm import tqdm
 
 from redbot import VersionInfo
 from redbot.core import data_manager
@@ -62,6 +61,8 @@
     "fetch_latest_red_version_info",
     "deprecated_removed",
     "RichIndefiniteBarColumn",
+    "RichSpeedColumn",
+    "detailed_progress",
     "cli_level_to_log_level",
 )
 
@@ -276,9 +277,10 @@ async def create_backup(dest: Path = Path.home()) -> Optional[Path]:
             to_backup.append(f)
 
     with tarfile.open(str(backup_fpath), "w:gz") as tar:
-        progress_bar = tqdm(to_backup, desc="Compressing data", unit=" files", dynamic_ncols=True)
-        for f in progress_bar:
-            tar.add(str(f), arcname=str(f.relative_to(data_path)), recursive=False)
+        with detailed_progress(unit="files") as progress:
+            progress_tracker = progress.track(to_backup, description="Compressing data")
+            for f in progress_tracker:
+                tar.add(str(f), arcname=str(f.relative_to(data_path)), recursive=False)
 
         # add repos backup
         repos_data = json.dumps(repo_output, indent=4)
@@ -403,9 +405,9 @@ def deprecated_removed(
     )
 
 
-class RichIndefiniteBarColumn(ProgressColumn):
-    def render(self, task):
-        return ProgressBar(
+class RichIndefiniteBarColumn(rich.progress.ProgressColumn):
+    def render(self, task: rich.progress.Task) -> rich.progress.ProgressBar:
+        return rich.progress.ProgressBar(
             pulse=task.completed < task.total,
             animation_time=task.get_time(),
             width=40,
@@ -414,6 +416,32 @@ def render(self, task):
         )
 
 
+class RichSpeedColumn(rich.progress.ProgressColumn):
+    def __init__(self, *, unit: str) -> None:
+        self.unit = unit
+        super().__init__()
+
+    def render(self, task: rich.progress.Task) -> Text:
+        speed = task.finished_speed or task.speed
+        if speed is None:
+            return Text("?", style="progress.data.speed")
+        return Text(f"{int(speed)} {self.unit}/s", style="progress.data.speed")
+
+
+def detailed_progress(*, unit: str) -> rich.progress.Progress:
+    return rich.progress.Progress(
+        rich.progress.SpinnerColumn(),
+        rich.progress.TextColumn("[progress.description]{task.description}"),
+        rich.progress.BarColumn(bar_width=None),
+        RichSpeedColumn(unit=unit),
+        rich.progress.TaskProgressColumn(),
+        rich.progress.TextColumn("eta"),
+        rich.progress.TimeRemainingColumn(),
+        rich.progress.TextColumn("elapsed"),
+        rich.progress.TimeElapsedColumn(),
+    )
+
+
 def cli_level_to_log_level(level: int) -> int:
     if level == 0:
         log_level = logging.INFO

redbot/setup.py

@@ -17,14 +17,14 @@
 from typing import Any, Dict, IO, List, Optional, Set, Tuple, Union
 
 import click
-from tqdm import tqdm
 
 from redbot.cogs.downloader.repo_manager import RepoManager
 from redbot.core._cli import confirm
 from redbot.core.utils._internal_utils import (
     safe_delete,
     create_backup as red_create_backup,
     cli_level_to_log_level,
+    detailed_progress,
 )
 from redbot.core import config, data_manager
 from redbot.core._config import migrate
@@ -466,8 +466,7 @@ def all_tar_members(self) -> List[tarfile.TarInfo]:
     def all_tar_member_names(self) -> List[str]:
         return [tarinfo.name for tarinfo in self.all_tar_members]
 
-    @functools.cached_property
-    def tar_members_to_extract(self) -> List[tarfile.TarInfo]:
+    def get_tar_members_to_extract(self) -> List[tarfile.TarInfo]:
         ignored_members: Set[str] = {"backup.version", "instance.json"}
         if not self.restore_downloader:
             ignored_members |= {
@@ -550,11 +549,11 @@ def _ask_for_storage(self) -> None:
         self.storage_details = driver_cls.get_config_details()
 
     def extractall(self) -> None:
-        progress_bar = tqdm(
-            self.tar_members_to_extract, desc="Extracting data", unit=" files", dynamic_ncols=True
-        )
-        # tar.errorlevel == 0 so errors are printed to stderr
-        self.tar.extractall(path=self.data_path, members=progress_bar)
+        to_extract = self.get_tar_members_to_extract()
+        with detailed_progress(unit="files") as progress:
+            progress_tracker = progress.track(to_extract, description="Extracting data")
+            # tar.errorlevel == 0 so errors are printed to stderr
+            self.tar.extractall(path=self.data_path, members=progress_tracker)
 
     def get_basic_config(self, use_json: bool = False) -> dict:
         default_dirs = deepcopy(data_manager.basic_config_default)