Force conversion unit in format_size

humanfriendly/__init__.py

@@ -157,7 +157,7 @@ def coerce_seconds(value):
     return value
 
 
-def format_size(num_bytes, keep_width=False, binary=False):
+def format_size(num_bytes, keep_width=False, binary=False, force_unit=None):
     """
     Format a byte count as a human readable file size.
 
@@ -166,6 +166,7 @@ def format_size(num_bytes, keep_width=False, binary=False):
                        :data:`False` if they can be stripped.
     :param binary: :data:`True` to use binary multiples of bytes (base-2),
                    :data:`False` to use decimal multiples of bytes (base-10).
+    :param force_unit: unit to force when converting the size
     :returns: The corresponding human readable file size (a string).
 
     This function knows how to format sizes in bytes, kilobytes, megabytes,
@@ -184,7 +185,48 @@ def format_size(num_bytes, keep_width=False, binary=False):
     '1 KiB'
     >>> format_size(1000 ** 3 * 4)
     '4 GB'
+    >>> format_size(1000 ** 3 * 4, force_unit='MB')
+    '4000 MB'
     """
+    if force_unit is not None:
+        force_unit_found = False
+        if binary:
+            symbols = [u.binary.symbol for u in disk_size_units]
+        else:
+            symbols = [u.decimal.symbol for u in disk_size_units]
+        force_unit_found = force_unit in symbols
+        if force_unit_found:
+            # this is going to fail for bytes
+            force_tuple_index = symbols.index(force_unit)
+            if binary:
+                number = round_number(
+                    float(num_bytes) /
+                    disk_size_units[force_tuple_index].binary.divider,
+                    keep_width=keep_width)
+                return pluralize(
+                    number,
+                    disk_size_units[force_tuple_index].binary.symbol,
+                    disk_size_units[force_tuple_index].binary.symbol)
+            else:
+                number = round_number(
+                    float(num_bytes) /
+                    disk_size_units[force_tuple_index].decimal.divider,
+                    keep_width=keep_width)
+                return pluralize(
+                    number,
+                    disk_size_units[force_tuple_index].decimal.symbol,
+                    disk_size_units[force_tuple_index].decimal.symbol)
+        else:
+            if force_unit == 'bytes':
+                return pluralize(
+                    round_number(
+                        float(num_bytes),
+                        keep_width=keep_width),
+                    'byte')
+            else:
+                raise InvalidSize(
+                    "You are trying to convert to an invalid unit: {}".format(force_unit))
+
     for unit in reversed(disk_size_units):
         if num_bytes >= unit.binary.divider and binary:
             number = round_number(float(num_bytes) / unit.binary.divider, keep_width=keep_width)

humanfriendly/tests.py

@@ -503,6 +503,22 @@ def test_format_size(self):
         self.assertEqual('1 EB', format_size(1000 ** 6))
         self.assertEqual('1 ZB', format_size(1000 ** 7))
         self.assertEqual('1 YB', format_size(1000 ** 8))
+        self.assertEqual('0 bytes', format_size(0, force_unit='bytes'))
+        self.assertEqual('1 byte', format_size(1, force_unit='bytes'))
+        self.assertEqual('42 bytes', format_size(42, force_unit='bytes'))
+        self.assertEqual('1 KB', format_size(1000 ** 1, force_unit='KB'))
+        self.assertEqual('1 MB', format_size(1000 ** 2, force_unit='MB'))
+        self.assertEqual('1 GB', format_size(1000 ** 3, force_unit='GB'))
+        self.assertEqual('1 TB', format_size(1000 ** 4, force_unit='TB'))
+        self.assertEqual('1 PB', format_size(1000 ** 5, force_unit='PB'))
+        self.assertEqual('1 EB', format_size(1000 ** 6, force_unit='EB'))
+        self.assertEqual('1 ZB', format_size(1000 ** 7, force_unit='ZB'))
+        self.assertEqual('1 YB', format_size(1000 ** 8, force_unit='YB'))
+        self.assertEqual('4 TB', format_size(1000 ** 4 * 4, force_unit='TB'))
+        self.assertEqual('4000 GB', format_size(1000 ** 4 * 4, force_unit='GB'))
+        self.assertEqual('4000000 MB', format_size(1000 ** 4 * 4, force_unit='MB'))
+        self.assertEqual('4000000000 KB', format_size(1000 ** 4 * 4, force_unit='KB'))
+        self.assertEqual('4000000000000 bytes', format_size(1000 ** 4 * 4, force_unit='bytes'))
         self.assertEqual('1 KiB', format_size(1024 ** 1, binary=True))
         self.assertEqual('1 MiB', format_size(1024 ** 2, binary=True))
         self.assertEqual('1 GiB', format_size(1024 ** 3, binary=True))
@@ -511,6 +527,19 @@ def test_format_size(self):
         self.assertEqual('1 EiB', format_size(1024 ** 6, binary=True))
         self.assertEqual('1 ZiB', format_size(1024 ** 7, binary=True))
         self.assertEqual('1 YiB', format_size(1024 ** 8, binary=True))
+        self.assertEqual('1 KiB', format_size(1024 ** 1, binary=True, force_unit='KiB'))
+        self.assertEqual('1 MiB', format_size(1024 ** 2, binary=True, force_unit='MiB'))
+        self.assertEqual('1 GiB', format_size(1024 ** 3, binary=True, force_unit='GiB'))
+        self.assertEqual('1 TiB', format_size(1024 ** 4, binary=True, force_unit='TiB'))
+        self.assertEqual('1 PiB', format_size(1024 ** 5, binary=True, force_unit='PiB'))
+        self.assertEqual('1 EiB', format_size(1024 ** 6, binary=True, force_unit='EiB'))
+        self.assertEqual('1 ZiB', format_size(1024 ** 7, binary=True, force_unit='ZiB'))
+        self.assertEqual('1 YiB', format_size(1024 ** 8, binary=True, force_unit='YiB'))
+        self.assertEqual('4 TiB', format_size(1024 ** 4 * 4, binary=True, force_unit='TiB'))
+        self.assertEqual('4096 GiB', format_size(1024 ** 4 * 4, binary=True, force_unit='GiB'))
+        self.assertEqual('4194304 MiB', format_size(1024 ** 4 * 4, binary=True, force_unit='MiB'))
+        self.assertEqual('4294967296 KiB', format_size(1024 ** 4 * 4, binary=True, force_unit='KiB'))
+        self.assertEqual('4398046511104 bytes', format_size(1024 ** 4 * 4, binary=True, force_unit='bytes'))
         self.assertEqual('45 KB', format_size(1000 * 45))
         self.assertEqual('2.9 TB', format_size(1000 ** 4 * 2.9))
 
@@ -787,8 +816,8 @@ def test_spinner(self):
                         .replace(ANSI_HIDE_CURSOR, ''))
         lines = [line for line in output.split(ANSI_ERASE_LINE) if line]
         self.assertTrue(len(lines) > 0)
-        self.assertTrue(all('test spinner' in l for l in lines))
-        self.assertTrue(all('%' in l for l in lines))
+        self.assertTrue(all('test spinner' in line for line in lines))
+        self.assertTrue(all('%' in line for line in lines))
         self.assertEqual(sorted(set(lines)), sorted(lines))
 
     def test_automatic_spinner(self):
@@ -952,7 +981,7 @@ def test_cli(self):
         # https://github.com/xolox/python-humanfriendly/issues/28
         returncode, output = run_cli(main, '--demo')
         assert returncode == 0
-        lines = [ansi_strip(l) for l in output.splitlines()]
+        lines = [ansi_strip(line) for line in output.splitlines()]
         assert "Text styles:" in lines
         assert "Foreground colors:" in lines
         assert "Background colors:" in lines

humanfriendly/usage.py

@@ -258,7 +258,7 @@ def render_usage(text):
                 ('\n\n'.join(render_paragraph(p, meta_variables) for p in split_paragraphs(description))).rstrip(),
             ])
         csv_lines = csv_buffer.getvalue().splitlines()
-        output.append('\n'.join('   %s' % l for l in csv_lines))
+        output.append('\n'.join('   %s' % line for line in csv_lines))
     logger.debug("Rendered output: %s", output)
     return '\n\n'.join(trim_empty_lines(o) for o in output)