[tor-commits] [stem/master] Add a 'round' argument to str_tools.size_label()

[atagar at torproject.org]

commit 0b1a0f947329d2b1a6272135800067ad07754071
Author: Damian Johnson <atagar at torproject.org>
Date:   Wed Sep 13 18:04:13 2017 -0700

    Add a 'round' argument to str_tools.size_label()
    
    Nyx is rounding down where I'd rather if it didn't. Adding an argument so we
    can customize this.
---
 docs/change_log.rst         |  1 +
 stem/util/str_tools.py      | 21 ++++++++++++++-------
 test/unit/util/str_tools.py |  5 +++++
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index 1bf01c5d..dafdb471 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -80,6 +80,7 @@ The following are only available within Stem's `git repository
   * Support connection resolution on OpenBSD using fstat (:trac:`13807`)
   * Added timeout argument to :func:`~stem.util.system.call`
   * Added cwd argument to :func:`~stem.util.system.call`
+  * Added round argument to :func:`~stem.util.str_tools.size_label`
   * Added :class:`~stem.util.test_tools.TimedTestRunner` and :func:`~stem.util.test_tools.test_runtimes`
   * Supporing pid arguments in :func:`~stem.util.system.is_running`
   * Normalized :func:`~stem.util.term.format` to return unicode
diff --git a/stem/util/str_tools.py b/stem/util/str_tools.py
index 696e6321..14f8a7cf 100644
--- a/stem/util/str_tools.py
+++ b/stem/util/str_tools.py
@@ -299,7 +299,7 @@ def crop(msg, size, min_word_length = 4, min_crop = 0, ending = Ending.ELLIPSE,
   return (return_msg, remainder) if get_remainder else return_msg
 
 
-def size_label(byte_count, decimal = 0, is_long = False, is_bytes = True):
+def size_label(byte_count, decimal = 0, is_long = False, is_bytes = True, round = False):
   """
   Converts a number of bytes into a human readable label in its most
   significant units. For instance, 7500 bytes would return "7 KB". If the
@@ -318,18 +318,22 @@ def size_label(byte_count, decimal = 0, is_long = False, is_bytes = True):
     >>> size_label(1050, 3, True)
     '1.025 Kilobytes'
 
+  .. versionchanged:: 1.6.0
+     Added round argument.
+
   :param int byte_count: number of bytes to be converted
   :param int decimal: number of decimal digits to be included
   :param bool is_long: expands units label
   :param bool is_bytes: provides units in bytes if **True**, bits otherwise
+  :param bool round: rounds normally if **True**, otherwise rounds down
 
   :returns: **str** with human readable representation of the size
   """
 
   if is_bytes:
-    return _get_label(SIZE_UNITS_BYTES, byte_count, decimal, is_long)
+    return _get_label(SIZE_UNITS_BYTES, byte_count, decimal, is_long, round)
   else:
-    return _get_label(SIZE_UNITS_BITS, byte_count, decimal, is_long)
+    return _get_label(SIZE_UNITS_BITS, byte_count, decimal, is_long, round)
 
 
 def time_label(seconds, decimal = 0, is_long = False):
@@ -542,7 +546,7 @@ def _parse_iso_timestamp(entry):
   return timestamp + datetime.timedelta(microseconds = int(microseconds))
 
 
-def _get_label(units, count, decimal, is_long):
+def _get_label(units, count, decimal, is_long, round = False):
   """
   Provides label corresponding to units of the highest significance in the
   provided set. This rounds down (ie, integer truncation after visible units).
@@ -552,6 +556,7 @@ def _get_label(units, count, decimal, is_long):
   :param int count: number of base units being converted
   :param int decimal: decimal precision of label
   :param bool is_long: uses the long label if **True**, short label otherwise
+  :param bool round: rounds normally if **True**, otherwise rounds down
   """
 
   # formatted string for the requested number of digits
@@ -566,10 +571,12 @@ def _get_label(units, count, decimal, is_long):
 
   for count_per_unit, short_label, long_label in units:
     if count >= count_per_unit:
-      # Rounding down with a '%f' is a little clunky. Reducing the count so
-      # it'll divide evenly as the rounded down value.
+      if not round:
+        # Rounding down with a '%f' is a little clunky. Reducing the count so
+        # it'll divide evenly as the rounded down value.
+
+        count -= count % (count_per_unit / (10 ** decimal))
 
-      count -= count % (count_per_unit / (10 ** decimal))
       count_label = label_format % (count / count_per_unit)
 
       if is_long:
diff --git a/test/unit/util/str_tools.py b/test/unit/util/str_tools.py
index 1576c06d..979f0896 100644
--- a/test/unit/util/str_tools.py
+++ b/test/unit/util/str_tools.py
@@ -67,6 +67,11 @@ class TestStrTools(unittest.TestCase):
     self.assertEqual('1.02 KB', str_tools.size_label(1050, 2))
     self.assertEqual('1.025 Kilobytes', str_tools.size_label(1050, 3, True))
 
+    self.assertEqual('1 KB', str_tools.size_label(1050, 0, round = True))
+    self.assertEqual('1.0 KB', str_tools.size_label(1050, 1, round = True))
+    self.assertEqual('1.03 KB', str_tools.size_label(1050, 2, round = True))
+    self.assertEqual('1.025 KB', str_tools.size_label(1050, 3, round = True))
+
     self.assertEqual('0 B', str_tools.size_label(0))
     self.assertEqual('0 Bytes', str_tools.size_label(0, is_long = True))
     self.assertEqual('0.00 B', str_tools.size_label(0, 2))