Unify and make symmetric pretty_size and parse_size

2 files changed, 24 insertions, 21 deletions

diff --git a/netlib/human.py b/netlib/human.py
index f4640c00..9eccd35b 100644
--- a/netlib/human.py
+++ b/netlib/human.py
@@ -1,26 +1,25 @@
-SIZE_UNITS = dict(
-    b=1024 ** 0,
-    k=1024 ** 1,
-    m=1024 ** 2,
-    g=1024 ** 3,
-    t=1024 ** 4,
-)
+
+SIZE_TABLE = [
+    ("b", 1024 ** 0),
+    ("k", 1024 ** 1),
+    ("m", 1024 ** 2),
+    ("g", 1024 ** 3),
+    ("t", 1024 ** 4),
+]
+
+SIZE_UNITS = dict(SIZE_TABLE)
 
 
 def pretty_size(size):
-    suffixes = [
-        ("B", 2 ** 10),
-        ("kB", 2 ** 20),
-        ("MB", 2 ** 30),
-    ]
-    for suf, lim in suffixes:
-        if size >= lim:
-            continue
-        else:
-            x = round(size / float(lim / 2 ** 10), 2)
+    for bottom, top in zip(SIZE_TABLE, SIZE_TABLE[1:]):
+        if bottom[1] <= size < top[1]:
+            suf = bottom[0]
+            lim = bottom[1]
+            x = round(size / lim, 2)
             if x == int(x):
                 x = int(x)
             return str(x) + suf
+    return "%s%s"%(size, SIZE_TABLE[0][0])
 
 
 def parse_size(s):
diff --git a/test/netlib/test_human.py b/test/netlib/test_human.py
index 3a445c0b..464d4646 100644
--- a/test/netlib/test_human.py
+++ b/test/netlib/test_human.py
@@ -1,6 +1,8 @@
 from netlib import human, tutils
 
 def test_parse_size():
+    assert human.parse_size("0") == 0
+    assert human.parse_size("0b") == 0
     assert human.parse_size("1") == 1
     assert human.parse_size("1k") == 1024
     assert human.parse_size("1m") == 1024**2
@@ -10,10 +12,12 @@ def test_parse_size():
 
 
 def test_pretty_size():
-    assert human.pretty_size(100) == "100B"
-    assert human.pretty_size(1024) == "1kB"
-    assert human.pretty_size(1024 + (1024 / 2.0)) == "1.5kB"
-    assert human.pretty_size(1024 * 1024) == "1MB"
+    assert human.pretty_size(0) == "0b"
+    assert human.pretty_size(100) == "100b"
+    assert human.pretty_size(1024) == "1k"
+    assert human.pretty_size(1024 + (1024 / 2.0)) == "1.5k"
+    assert human.pretty_size(1024 * 1024) == "1m"
+    assert human.pretty_size(10 * 1024 * 1024) == "10m"
 
 
 def test_pretty_duration():