Fix "Nines" bug and add tests

millify/__init__.py

@@ -1,38 +1,37 @@
-import math
 import re
-from decimal import Decimal
 
 __author__ = "Alexander Zaitsev (azaitsev@gmail.com)"
 __copyright__ = "Copyright 2018, azaitsev@gmail.com"
 __license__ = "MIT"
 __version__ = "0.1.1"
 
 
-def remove_exponent(d):
-    """Remove exponent."""
-    return d.quantize(Decimal(1)) if d == d.to_integral() else d.normalize()
-
-
 def millify(n, precision=0, drop_nulls=True, prefixes=[]):
     """Humanize number."""
     millnames = ['', 'k', 'M', 'B', 'T', 'P', 'E', 'Z', 'Y']
-    if prefixes:
+    if prefixes and isinstance(prefixes, list):
         millnames = ['']
         millnames.extend(prefixes)
     n = float(n)
-    millidx = max(0, min(len(millnames) - 1,
-                         int(math.floor(0 if n == 0 else math.log10(abs(n)) / 3))))
-    result = '{:.{precision}f}'.format(n / 10**(3 * millidx), precision=precision)
+    millidx = 0
+    while abs(n) >= 1000:
+        millidx += 1
+        n = round(n / 1000.0, precision)
+    if n < 1000:
+        n = round(n, precision)
+    result = '{}'.format(n)
     if drop_nulls:
-        result = remove_exponent(Decimal(result))
-    return '{0}{dx}'.format(result, dx=millnames[millidx])
+        result = result.rstrip('0').rstrip('.')
+    return '{}{dx}'.format(result, dx=millnames[millidx])
 
 
 def prettify(amount, separator=','):
     """Separate with predefined separator."""
+    if not isinstance(separator, str):
+        separator = ','
     orig = str(amount)
     new = re.sub("^(-?\d+)(\d{3})", "\g<1>{0}\g<2>".format(separator), str(amount))
     if orig == new:
         return new
     else:
-        return prettify(new)
+        return prettify(new, separator=separator)

tests/test_millify.py

@@ -0,0 +1,82 @@
+import unittest
+from millify import millify
+from millify import prettify
+
+
+class TestMillify(unittest.TestCase):
+    """ Tests the millify() function. """
+
+    def test_basics(self):
+        self.assertEqual(millify(123), '123')
+        self.assertEqual(millify(1234), '1k')
+        self.assertEqual(millify(12345), '12k')
+        self.assertEqual(millify(12345678), '12M')
+        self.assertEqual(millify(1000000000), '1B')
+
+    def test_precision(self):
+        self.assertEqual(millify(1234, precision=2), '1.23k')
+        self.assertEqual(millify(12345, precision=3), '12.345k')
+        self.assertEqual(millify(12342123, precision=2), '12.34M')
+        self.assertEqual(millify(12345123, precision=2), '12.35M')
+        self.assertEqual(millify(12345678, precision=3), '12.346M')
+        self.assertEqual(millify(12345678, precision=4), '12.3457M')
+
+    def test_strings(self):
+        self.assertEqual(millify('1234'), '1k')
+        self.assertEqual(millify('12345'), '12k')
+        self.assertEqual(millify('12345678'), '12M')
+
+    def test_drop_nulls(self):
+        self.assertEqual(millify(10000, precision=1, drop_nulls=False), '10.0k')
+        self.assertEqual(millify(10000, precision=1, drop_nulls=True), '10k')
+        self.assertEqual(millify(10000, precision=1), '10k')
+
+    def test_prefixes(self):
+        prefixes = ['kB', 'MB', 'GB']
+        self.assertEqual(millify(10000, prefixes=None), '10k')
+        self.assertEqual(millify(10000, prefixes=True), '10k')
+        self.assertEqual(millify(10000, prefixes=prefixes), '10kB')
+        self.assertEqual(millify(1000000, prefixes=prefixes), '1MB')
+        self.assertEqual(millify(1000000000, prefixes=prefixes), '1GB')
+
+    def test_nines(self):
+        self.assertEqual(millify(999), '999')
+        self.assertEqual(millify('999'), '999')
+        self.assertEqual(millify(9999), '10k')
+        self.assertEqual(millify(9999, precision=2), '10k')
+        self.assertEqual(millify(99999), '100k')
+        self.assertEqual(millify(999999), '1M')
+        self.assertEqual(millify(999999999), '1B')
+
+
+class TestPrettify(unittest.TestCase):
+    """ Tests the prettify() function. """
+
+    def test_basics(self):
+        self.assertEqual(prettify(1234), '1,234')
+        self.assertEqual(prettify(12345), '12,345')
+        self.assertEqual(prettify(123456), '123,456')
+        self.assertEqual(prettify(1234567), '1,234,567')
+        self.assertEqual(prettify(1234567890), '1,234,567,890')
+
+    def test_strings(self):
+        self.assertEqual(prettify('1234'), '1,234')
+        self.assertEqual(prettify('12345'), '12,345')
+        self.assertEqual(prettify('123456'), '123,456')
+        self.assertEqual(prettify('1234567'), '1,234,567')
+        self.assertEqual(prettify('1234567890'), '1,234,567,890')
+
+    def test_custom_separator(self):
+        self.assertEqual(prettify('1234', separator=','), '1,234')
+        self.assertEqual(prettify('12345', separator='.'), '12.345')
+        self.assertEqual(prettify('123456', separator='/'), '123/456')
+        self.assertEqual(prettify('1234567', separator='-'), '1-234-567')
+        self.assertEqual(prettify('1234567890', separator=' '), '1 234 567 890')
+
+    def test_separator_type(self):
+        self.assertEqual(prettify('1234567890', separator=None), '1,234,567,890')
+        self.assertEqual(prettify('1234567890', separator=True), '1,234,567,890')
+
+
+if __name__ == '__main__':
+    unittest.main()