chore: include more detailed steps and change ways grouping

kraina-ai · Dec 30, 2023 · 931421f · 931421f
1 parent 37c5ad8
commit 931421f
Show file tree

Hide file tree

Showing 3 changed files with 98 additions and 64 deletions.
diff --git a/pdm.lock b/pdm.lock
diff --git a/quackosm/_rich_progress.py b/quackosm/_rich_progress.py
@@ -5,6 +5,7 @@
 
 __all__ = ["TaskProgressSpinner", "TaskProgressBar"]
 
+TOTAL_STEPS = 34
 
 class TaskProgressSpinner:
     def __init__(self, step_name: str, step_number: str):
@@ -18,7 +19,7 @@ def __enter__(self):
 
             self.progress = Progress(
                 SpinnerColumn(),
-                TextColumn(f"[{self.step_number: >4}/18]"),
+                TextColumn(f"[{self.step_number: >4}/{TOTAL_STEPS}]"),
                 TextColumn("[progress.description]{task.description}"),
                 TextColumn("•"),
                 TimeElapsedColumn(),
@@ -58,7 +59,7 @@ def __enter__(self):
 
             self.progress = Progress(
                 SpinnerColumn(),
-                TextColumn(f"[{self.step_number: >4}/18]"),
+                TextColumn(f"[{self.step_number: >4}/{TOTAL_STEPS}]"),
                 TextColumn(
                     "[progress.description]{task.description}"
                     " [progress.percentage]{task.percentage:>3.0f}%"
@@ -87,7 +88,7 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
     def track(self, iterable: Iterable):
         if self.progress is not None:
-            for i in self.progress.track(iterable, description=self.step_name):
+            for i in self.progress.track(list(iterable), description=self.step_name):
                 yield i
         else:
             for i in iterable: