eng_str(): Adding unit argument.

2 files changed, 58 insertions, 16 deletions

diff --git a/src/ee/formatting/__init__.py b/src/ee/formatting/__init__.py
index a1e27a4..42bf48d 100644
--- a/src/ee/formatting/__init__.py
+++ b/src/ee/formatting/__init__.py
@@ -59,18 +59,40 @@ def e_series_find_closest(value):
 
   return min(series, key=lambda v: abs(v - value)) * 10**e
 
-def eng_str(value):
-  big = ['', ' k', ' M', ' G', ' T']
-  e_floor = math.floor(math.log10(value))
+import numpy
+def eng_str(value, unit = None):
 
-  div3 = int(math.floor(e_floor/3))
-  suffix = big[div3]
-  scale = 10**(div3*3)
-  scaled = value/scale
-
-  if scaled == int(scaled):
-    return "{:1.0f}{}".format(scaled, suffix)
+  if value == 0:
+    s = '0' + (' ' + unit if unit is not None else '')
   else:
-    return "{:1.1f}{}".format(scaled, suffix)
+    big = ['k', 'M', 'G', 'T']
+    small = ['m', 'u', 'p']
+
+    if unit is not None:
+      big = [' ' + unit] + [' ' + x + unit for x in big]
+      small = [' ' + unit] + [' ' + x + unit for x in small]
+    else:
+      big = [''] + [' ' + x for x in big]
+      small = [''] + [' ' + x for x in small]
+
+    e_floor = math.floor(math.log10(value))
+
+    div3 = math.floor(e_floor/3)
+    if div3 > 0:
+      suffixes = big
+      idx = int(div3)
+    else:
+      suffixes = small
+      idx = int(-div3)
+
+    suffix = suffixes[idx]
+    scale = 10**(div3*3)
+    scaled = value/scale
+
+    if (scaled - math.floor(scaled)) < 0.0001:
+#    if scaled == math.floor(scaled):
+      s = "{:1.0f}{}".format(scaled, suffix)
+    else:
+      s = "{:1.1f}{}".format(scaled, suffix)
 
-#    return "scaled={:>9}, div3={}, floor(e)={}, scale={}, suffix={}".format(str(scaled), div3, e_floor, scale, suffix)
+  return s
diff --git a/test/test_formatting.py b/test/test_formatting.py
index 560daee..6fa0b07 100644
--- a/test/test_formatting.py
+++ b/test/test_formatting.py
@@ -1,7 +1,9 @@
 import pytest
+import numpy as np
 from ee.formatting import eng_str
 
 @pytest.mark.parametrize("input,expected", [
+    (0, "0"),
     (5.5, "5.5"),
     (55, "55"),
     (550, "550"),
@@ -9,9 +11,27 @@ from ee.formatting import eng_str
     (55000, "55 k"),
     (550000, "550 k"),
     (5500000, "5.5 M"),
-    ])
+    (0.5, "500 m"),
+    (0.055, "55 m"),
+    (0.0055, "5.5 m"),
+    (0.00055, "550 u"),
+    (0.000055, "55 u"),
+    (0.0000055, "5.5 u"),
+    (0.00000055, "550 p"),
+  ])
 def test_eng_str(input, expected):
-  assert eng_str(input) == expected
+  assert expected == eng_str(input)
+  npinput = np.float64(input)
+  assert expected == eng_str(npinput)
+
+@pytest.mark.parametrize("input,expected", [
+    (0, "0 A"),
+    (5.5, "5.5 A"),
+    (1100, "1.1 kA"),
+    (0.05, "50 mA"),
+  ])
+def test_eng_str_with_unit(input, expected):
+  assert expected == eng_str(input, unit = 'A')
 
 @pytest.mark.parametrize("input,expected", [
     (100, ''), (101, ''), (102, ''), (103, ''), (104, ''), (105, ''), (106, ''), (107, ''), (108, ''), (109, ''),
@@ -20,7 +40,7 @@ def test_eng_str(input, expected):
     (130, ''), (131, ''), (132, ''), (133, ''), (134, ''), (135, ''), (136, ''), (137, ''), (138, ''), (139, ''),
   ])
 def xx_test_eng_str2(input, expected):
-  assert eng_str(input) == expected
+  assert expected == eng_str(input)
 
 @pytest.mark.parametrize("input,expected", [
     (10, '10'),
@@ -35,4 +55,4 @@ def xx_test_eng_str2(input, expected):
     (19, '19'),
   ])
 def xx_test_eng_str3(input, expected):
-  assert eng_str(input) == expected
+  assert expected == eng_str(input)