Move new behaviors to TimeRemainingColumn

rich/progress.py

@@ -343,40 +343,46 @@ def render(self, task: "Task") -> Text:
 
 
 class TimeRemainingColumn(ProgressColumn):
-    """Renders estimated time remaining."""
+    """Renders estimated time remaining.
+
+    Args:
+        compact (bool, optional): Render MM:SS when time remaining is less than an hour. Defaults to False.
+        elapsed_when_finished (bool, optional): Render time elapsed when the task is finished. Defaults to False.
+    """
 
     # Only refresh twice a second to prevent jitter
     max_refresh = 0.5
 
+    def __init__(
+        self,
+        compact: bool = False,
+        elapsed_when_finished: bool = False,
+        table_column: Optional[Column] = None,
+    ):
+        self.compact = compact
+        self.elapsed_when_finished = elapsed_when_finished
+        super().__init__(table_column=table_column)
+
     def render(self, task: "Task") -> Text:
         """Show time remaining."""
-        remaining = task.time_remaining
-        if remaining is None:
-            return Text("-:--:--", style="progress.remaining")
-        remaining_delta = timedelta(seconds=int(remaining))
-        return Text(str(remaining_delta), style="progress.remaining")
-
-
-class CondensedTimeColumn(ProgressColumn):
-    """Renders estimated time remaining, or elapsed time when the task is finished."""
-
-    # Only refresh twice a second to prevent jitter
-    max_refresh = 0.5
+        if self.elapsed_when_finished and task.finished:
+            task_time = task.finished_time
+            style = "progress.elapsed"
+        else:
+            task_time = task.time_remaining
+            style = "progress.remaining"
 
-    def render(self, task: "Task") -> Text:
-        """Show time."""
-        style = "progress.elapsed" if task.finished else "progress.remaining"
-        task_time = task.finished_time if task.finished else task.time_remaining
         if task_time is None:
-            return Text("--:--", style=style)
+            return Text("--:--" if self.compact else "-:--:--", style=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 hours:
-            formatted = f"{hours:d}:{minutes:02d}:{seconds:02d}"
-        else:
+
+        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=style)
 

tests/test_progress.py

@@ -23,7 +23,6 @@
     TextColumn,
     TimeElapsedColumn,
     TimeRemainingColumn,
-    CondensedTimeColumn,
     track,
     _TrackThread,
     TaskID,
@@ -91,7 +90,6 @@ class FakeTask(Task):
     assert str(text) == "0:01:00"
 
 
-@pytest.mark.parametrize("finished", [False, True])
 @pytest.mark.parametrize(
     "task_time, formatted",
     [
@@ -102,13 +100,20 @@ class FakeTask(Task):
         (4210, "1:10:10"),
     ],
 )
-def test_condensed_time_column(finished, task_time, formatted):
-    if finished:
-        task = SimpleNamespace(finished=finished, finished_time=task_time)
-    else:
-        task = SimpleNamespace(finished=finished, time_remaining=task_time)
+def test_compact_time_remaining_column(task_time, formatted):
+    task = SimpleNamespace(finished=False, time_remaining=task_time)
+    column = TimeRemainingColumn(compact=True)
+
+    assert str(column.render(task)) == formatted
+
+
+def test_time_remaining_column_elapsed_when_finished():
+    task_time = 71
+    formatted = "0:01:11"
+
+    task = SimpleNamespace(finished=True, finished_time=task_time)
+    column = TimeRemainingColumn(elapsed_when_finished=True)
 
-    column = CondensedTimeColumn()
     assert str(column.render(task)) == formatted