Skip to content

Commit

Permalink
Add flag to render in binary/decimal suffix in DownloadColumn
Browse files Browse the repository at this point in the history
 Add a test for rendering binary suffix in test_progress.py
  • Loading branch information
amartya-dev committed Oct 13, 2020
1 parent d17b7f5 commit 630ed8b
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 5 deletions.
24 changes: 19 additions & 5 deletions rich/progress.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,6 @@

ProgressType = TypeVar("ProgressType")


GetTimeCallable = Callable[[], float]


Expand Down Expand Up @@ -292,15 +291,30 @@ def render(self, task: "Task") -> Text:


class DownloadColumn(ProgressColumn):
"""Renders file size downloaded and total, e.g. '0.5/2.3 GB'."""
"""Renders file size downloaded and total, e.g. '0.5/2.3 GB'.
Args:
decimal_suffix (bool, optional): Flag to renser filesize in desired format, disable to render in binary. Defaults to True.
"""

def __init__(self, decimal_suffix: bool = True) -> None:
self.decimal_ssuffix = decimal_suffix
super().__init__()

def render(self, task: "Task") -> Text:
"""Calculate common unit for completed and total."""
completed = int(task.completed)
total = int(task.total)
unit, suffix = filesize.pick_unit_and_suffix(
total, ["bytes", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"], 1000
)
if self.decimal_ssuffix:
unit, suffix = filesize.pick_unit_and_suffix(
total, ["bytes", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB"], 1000
)
else:
unit, suffix = filesize.pick_unit_and_suffix(
total,
["bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"],
1024,
)
completed_ratio = completed / unit
total_ratio = total / unit
precision = 0 if unit == 1 else 1
Expand Down
9 changes: 9 additions & 0 deletions tests/test_progress.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,15 @@ def test_download_progress_uses_decimal_units() -> None:
assert rendered_progress == expected


def test_download_progress_uses_binary_units() -> None:

column = DownloadColumn(decimal_suffix=False)
test_task = Task(1, "test", 1024, 512, _get_time=lambda: 1.0)
rendered_progress = str(column.render(test_task))
expected = "0.5/1.0 KiB"
assert rendered_progress == expected


def test_task_ids():
progress = make_progress()
assert progress.task_ids == [0, 1, 2, 4]
Expand Down

0 comments on commit 630ed8b

Please sign in to comment.