Merge #8: Allow correct behavior of 'parse_size' and 'format_size'

Author: xolox

humanfriendly/__init__.py

@@ -59,6 +59,14 @@
                    dict(prefix='t', divider=1024**4, singular='TB', plural='TB'),
                    dict(prefix='p', divider=1024**5, singular='PB', plural='PB'))
 
+# Common disk size units based on IEEE 1541.
+disk_size_units_ieee = (dict(prefix='b', divider=1, singular='byte', plural='bytes'),
+                        dict(prefix='k', divider=1000**1, singular='KB', plural='KB'),
+                        dict(prefix='m', divider=1000**2, singular='MB', plural='MB'),
+                        dict(prefix='g', divider=1000**3, singular='GB', plural='GB'),
+                        dict(prefix='t', divider=1000**4, singular='TB', plural='TB'),
+                        dict(prefix='p', divider=1000**5, singular='PB', plural='PB'))
+
 # Common length size units, used for formatting and parsing.
 length_size_units = (dict(prefix='nm', divider=1e-09, singular='nm', plural='nm'),
                      dict(prefix='mm', divider=1e-03, singular='mm', plural='mm'),
@@ -104,7 +112,7 @@ def coerce_boolean(value):
         return bool(value)
 
 
-def format_size(num_bytes, keep_width=False):
+def format_size(num_bytes, keep_width=False, correct=False):
     """
     Format a byte count as a human readable file size.
 
@@ -128,15 +136,20 @@ def format_size(num_bytes, keep_width=False):
     '1 MB'
     >>> format_size(1024 ** 3 * 4)
     '4 GB'
+    >>> format_size(1000 ** 3 * 4, correct=True)
+    '4 GB'
     """
-    for unit in reversed(disk_size_units):
+    units = disk_size_units
+    if correct:
+        units = disk_size_units_ieee
+    for unit in reversed(units):
         if num_bytes >= unit['divider']:
             number = round_number(float(num_bytes) / unit['divider'], keep_width=keep_width)
             return pluralize(number, unit['singular'], unit['plural'])
     return pluralize(num_bytes, 'byte')
 
 
-def parse_size(size):
+def parse_size(size, correct=False):
     """
     Parse a human readable data size and return the number of bytes.
 
@@ -155,6 +168,8 @@ def parse_size(size):
     5120
     >>> parse_size('1.5 GB')
     1610612736
+    >>> parse_size('1.5 GB', correct=True)
+    1500000000
     """
     tokens = tokenize(size)
     if tokens and isinstance(tokens[0], numbers.Number):
@@ -165,7 +180,10 @@ def parse_size(size):
         if len(tokens) == 2 and is_string(tokens[1]):
             normalized_unit = tokens[1].lower()
             # Try to match the first letter of the unit.
-            for unit in disk_size_units:
+            units = disk_size_units
+            if correct:
+                units = disk_size_units_ieee
+            for unit in units:
                 if normalized_unit.startswith(unit['prefix']):
                     return int(tokens[0] * unit['divider'])
     # We failed to parse the size specification.

humanfriendly/tests.py

@@ -190,6 +190,10 @@ def test_format_size(self):
         self.assertEqual('1 GB', humanfriendly.format_size(1024 ** 3))
         self.assertEqual('1 TB', humanfriendly.format_size(1024 ** 4))
         self.assertEqual('1 PB', humanfriendly.format_size(1024 ** 5))
+        self.assertEqual('1 byte', humanfriendly.format_size(1, correct=True))
+        self.assertEqual('45 KB', humanfriendly.format_size(1000 * 45, correct=True))
+        self.assertEqual('1 GB', humanfriendly.format_size(1000 ** 3, correct=True))
+        self.assertEqual('2.9 TB', humanfriendly.format_size(1000 ** 4 * 2.9, correct=True))
 
     def test_parse_size(self):
         """Test :func:`humanfriendly.parse_size()`."""
@@ -203,6 +207,8 @@ def test_parse_size(self):
         self.assertEqual(1024 ** 3 * 1.5, humanfriendly.parse_size('1.5 GB'))
         self.assertRaises(humanfriendly.InvalidSize, humanfriendly.parse_size, '1z')
         self.assertRaises(humanfriendly.InvalidSize, humanfriendly.parse_size, 'a')
+        self.assertEqual(1000, humanfriendly.parse_size('1 KB', correct=True))
+        self.assertEqual(1000 ** 2 * 69, humanfriendly.parse_size('69 MB', correct=True))
 
     def test_format_length(self):
         """Test :func:`humanfriendly.format_length()`."""