From 2c9e56e830d0a919a7fff0996b7ca677501a5273 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sat, 8 Oct 2022 16:32:15 -0700
Subject: [PATCH] Improve device size reporting in the description

i.e. support 512GB SSDs, 1.5TB HDDs, etc
Addresses issue #199
---
 scripts/wk/hw/disk.py | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py
index a663cd6b..0b4bec48 100644
--- a/scripts/wk/hw/disk.py
+++ b/scripts/wk/hw/disk.py
@@ -20,7 +20,7 @@ from wk.hw.smart import (
   get_known_disk_attributes,
   update_smart_details,
   )
-from wk.std import PLATFORM, bytes_to_string, color_string, strip_colors
+from wk.std import PLATFORM, color_string, strip_colors
 
 
 # STATIC VARIABLES
@@ -39,6 +39,7 @@ class Disk:
   attributes: dict[Any, dict] =       field(init=False, default_factory=dict)
   bus: str =                          field(init=False)
   children: list[dict] =              field(init=False, default_factory=list)
+  description: str =                  field(init=False)
   filesystem: str =                   field(init=False)
   known_attributes: dict[Any, dict] = field(init=False, default_factory=dict)
   log_sec: int =                      field(init=False)
@@ -59,6 +60,7 @@ class Disk:
   def __post_init__(self) -> None:
     self.path = pathlib.Path(self.path).resolve()
     self.update_details()
+    self.set_description()
     self.known_attributes = get_known_disk_attributes(self.model)
     enable_smart(self)
     update_smart_details(self)
@@ -88,14 +90,6 @@ class Disk:
         present = True
     return present
 
-  @property
-  def description(self) -> str:
-    """Get disk description from details."""
-    return (
-      f'{bytes_to_string(self.size, use_binary=False)}'
-      f' ({self.bus}) {self.model} {self.serial}'
-      )
-
   def disable_disk_tests(self) -> None:
     """Disable all tests."""
     LOG.warning('Disabling all tests for: %s', self.path)
@@ -162,6 +156,33 @@ class Disk:
       return False
     return True
 
+  def set_description(self) -> None:
+    """Set disk description from details."""
+    decimals = 1
+    suffix = ' '
+
+    # Set size_str (try binary scale first)
+    for scale in (1024, 1000):
+      size = float(self.size)
+      units = list('KMGTPEZY')
+      while units:
+        if abs(size) < scale:
+          break
+        size /= scale
+        suffix = units.pop(0)
+      size = ((size * 10) // 1) / 10
+      if size % 1 == 0:
+        # Found an exact whole number, drop the decimal
+        decimals = 0
+        break
+      if size % 1 == 0.5:
+        break
+
+    # Done
+    self.description = (
+      f'{size:0.{decimals}f} {suffix}B ({self.bus}) {self.model} {self.serial}'
+      )
+
   def update_details(self, skip_children=True) -> None:
     """Update disk details using OS specific methods.