diff --git a/humanfriendly/__init__.py b/humanfriendly/__init__.py index 3982aad..1a6c2ac 100644 --- a/humanfriendly/__init__.py +++ b/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) diff --git a/humanfriendly/tests.py b/humanfriendly/tests.py index 4e8c019..878dc28 100644 --- a/humanfriendly/tests.py +++ b/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 diff --git a/humanfriendly/usage.py b/humanfriendly/usage.py index 99a6b59..a045acb 100644 --- a/humanfriendly/usage.py +++ b/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)