refactor: tqdm -> rich.progress

scripts/download_tools.py

@@ -1,7 +1,28 @@
-import requests
 import os
-from tqdm import tqdm
-from config import config, log
+import requests
+import logging
+from config import config, log, console
+from rich.progress import (
+    BarColumn,
+    DownloadColumn,
+    Progress,
+    TextColumn,
+    TimeRemainingColumn,
+    TransferSpeedColumn,
+)
+
+progress = Progress(
+    TextColumn("[bold blue]{task.fields[filename]}", justify="right"),
+    BarColumn(bar_width=None),
+    "[progress.percentage]{task.percentage:>3.1f}%",
+    "•",
+    DownloadColumn(),
+    "•",
+    TransferSpeedColumn(),
+    "•",
+    TimeRemainingColumn(),
+    console=console
+)
 
 
 def check_if_tool_exists(tool: str) -> bool:
@@ -19,27 +40,30 @@ def check_if_tool_exists(tool: str) -> bool:
     else:
         return True
 
+requests_log = logging.getLogger("urllib3.connectionpool")
+requests_log.setLevel(logging.WARNING)
 
 def download_tool(url: str, tool: str):
     if not check_if_tool_exists(tool):
-        log.info(f"downloading a tool: `{tool}`")
+        progress.start()
+        
         try:
+            log.info(f"Requesting {url}")
             response = requests.get(url, stream=True)
-            total = int(response.headers.get("content-length", 0))
-            with open(f"{config['folders']['tools']}/{tool}", "wb") as file, tqdm(
-                desc=tool,
-                total=total,
-                unit="iB",
-                unit_scale=True,
-                unit_divisor=1024,
-            ) as bar:
-                for bytes in response.iter_content(chunk_size=8192):
+            total = int(response.headers.get("content-length", None))
+            task_id = progress.add_task("download", start=False, total=total, filename=tool)
+
+            with open(f"{config['folders']['tools']}/{tool}", "wb") as file:
+                progress.start_task(task_id)
+                for bytes in response.iter_content(chunk_size=32768):
                     size = file.write(bytes)
-                    bar.update(size)
+                    progress.update(task_id, advance=size)
+
             log.info(f"`{tool}` downloaded")
         except Exception as e:
             log.error(f"error while downloading `{tool}`: {e}")
 
+        progress.stop()
 
 def check_and_download_all_tools():
     for tool in config["tools"]:

scripts/select_patches.py

@@ -1,8 +1,10 @@
 import os, json
 import importlib
 from typing import TypedDict
+
 from beaupy import select_multiple
-from tqdm import tqdm
+from rich.progress import BarColumn, Progress, TextColumn
+
 from config import config, log, console
 
 
@@ -59,6 +61,14 @@ class PatchStatus(TypedDict):
     status: bool
 
 
+progress = Progress(
+    "[progress.description]{task.description}",
+    TextColumn(text_format="{task.fields[patch]}"),
+    BarColumn(bar_width=None),
+    "[blue]{task.completed}/{task.total}",
+)
+
+
 def apply_patches(patches: list[str]) -> list[PatchStatus]:
     modules = []
     statuses = []
@@ -69,25 +79,27 @@ def apply_patches(patches: list[str]) -> list[PatchStatus]:
         )
         modules.append(Patch(name, module))
     modules.sort(key=lambda x: x.package.priority, reverse=True)
-    
-    with tqdm(
-        total=len(modules),
-        unit="patch",
-        unit_divisor=1,
-    ) as bar:
-        for patch in modules:
-            bar.set_description(f"{patch.name}")
-            conf = {}
-            if os.path.exists(f"{config['folders']['patches']}/{patch.name}.config.json"):
+
+    with progress:
+        task = progress.add_task("applying patch:", total=len(modules), patch="")
+        for module in modules:
+            progress.update(task, patch=module.name)
+
+            patch_conf = {}
+            if os.path.exists(
+                f"{config['folders']['patches']}/{module.name}.config.json"
+            ):
                 with open(
-                    f"{config['folders']['patches']}/{patch.name}.config.json",
+                    f"{config['folders']['patches']}/{module.name}.config.json",
                     "r",
                     encoding="utf-8",
-                ) as conf:
-                    conf = json.loads(conf.read())
-            conf["src"] = config["folders"]["decompiled"]
-            status = patch.apply(conf)
-            statuses.append({"name": patch.name, "status": status})
-            bar.update()
+                ) as f:
+                    patch_conf = json.loads(f.read())
+
+            status = module.apply(patch_conf)
+            statuses.append({"name": module.name, "status": status})
+            progress.update(task, advance=1)
+
+        progress.update(task, description="patches applied", patch="")
 
     return statuses