Simplified custom column classes

nncf/common/logging/track_progress.py

@@ -20,22 +20,20 @@
 from rich.progress import Task
 from rich.progress import TaskProgressColumn
 from rich.progress import TextColumn
+from rich.progress import TimeElapsedColumn
+from rich.progress import TimeRemainingColumn
 from rich.style import StyleType
 from rich.text import Text
 
-INTEL_BLUE_COLOR = (0, 113, 197)
+INTEL_BLUE_COLOR = "#0068b5"
 
 
 class IterationsColumn(ProgressColumn):
-    def __init__(self, style: Union[str, StyleType]):
-        super().__init__()
-        self.style = style
-
     def render(self, task: Task) -> Text:
         if task.total is None:
             return Text("")
         text = f"{int(task.completed)}/{int(task.total)}"
-        return Text(text, style=self.style)
+        return Text(text, style=INTEL_BLUE_COLOR)
 
 
 class SeparatorColumn(ProgressColumn):
@@ -49,71 +47,16 @@ def render(self, task: Task) -> Text:
         return Text("•")
 
 
-class TimeElapsedColumnWithStyle(ProgressColumn):
-    """
-    Renders time elapsed.
-
-    Similar to TimeElapsedColumn, but with addition of style parameter.
-    """
-
-    def __init__(self, style: Union[str, StyleType]):
-        super().__init__()
-        self.style = style
-
+class TimeElapsedColumnWithStyle(TimeElapsedColumn):
     def render(self, task: "Task") -> Text:
-        """Show time elapsed."""
-        elapsed = task.finished_time if task.finished else task.elapsed
-        if elapsed is None:
-            return Text("-:--:--", style=self.style)
-        delta = timedelta(seconds=max(0, int(elapsed)))
-        return Text(str(delta), style=self.style)
-
+        text = super().render(task)
+        return Text(text._text[0], style=INTEL_BLUE_COLOR)
 
-class TimeRemainingColumnWithStyle(ProgressColumn):
-    """
-    Renders estimated time remaining.
-
-    Similar to TimeRemainingColumn, but with addition of style parameter.
-    """
-
-    # Only refresh twice a second to prevent jitter
-    max_refresh = 0.5
-
-    def __init__(
-        self,
-        style: Union[str, StyleType],
-        compact: bool = False,
-        elapsed_when_finished: bool = False,
-        table_column: Optional[Column] = None,
-    ):
-        self.style = style
-        self.compact = compact
-        self.elapsed_when_finished = elapsed_when_finished
-        super().__init__(table_column=table_column)
 
+class TimeRemainingColumnWithStyle(TimeRemainingColumn):
     def render(self, task: "Task") -> Text:
-        """Show time remaining."""
-        if self.elapsed_when_finished and task.finished:
-            task_time = task.finished_time
-        else:
-            task_time = task.time_remaining
-
-        if task.total is None:
-            return Text("", style=self.style)
-
-        if task_time is None:
-            return Text("--:--" if self.compact else "-:--:--", style=self.style)
-
-        # Based on https://github.com/tqdm/tqdm/blob/master/tqdm/std.py
-        minutes, seconds = divmod(int(task_time), 60)
-        hours, minutes = divmod(minutes, 60)
-
-        if self.compact and not hours:
-            formatted = f"{minutes:02d}:{seconds:02d}"
-        else:
-            formatted = f"{hours:d}:{minutes:02d}:{seconds:02d}"
-
-        return Text(formatted, style=self.style)
+        text = super().render(task)
+        return Text(text._text[0], style=INTEL_BLUE_COLOR)
 
 
 class track:
@@ -164,8 +107,6 @@ def __init__(
         self.update_period = update_period
         self.task = None
 
-        text_style = f"rgb({INTEL_BLUE_COLOR[0]},{INTEL_BLUE_COLOR[1]},{INTEL_BLUE_COLOR[2]})"
-
         self.columns: List[ProgressColumn] = (
             [TextColumn("[progress.description]{task.description}")] if description else []
         )
@@ -179,11 +120,11 @@ def __init__(
                     bar_width=None,
                 ),
                 TaskProgressColumn(show_speed=show_speed),
-                IterationsColumn(style=text_style),
+                IterationsColumn(),
                 SeparatorColumn(),
-                TimeElapsedColumnWithStyle(style=text_style),
+                TimeElapsedColumnWithStyle(),
                 SeparatorColumn(disable_if_no_total=True),  # disable because time remaining will be empty
-                TimeRemainingColumnWithStyle(style=text_style),
+                TimeRemainingColumnWithStyle(),
             )
         )