diff --git a/pydu/unit.py b/pydu/unit.py new file mode 100644 index 0000000..851a523 --- /dev/null +++ b/pydu/unit.py @@ -0,0 +1,31 @@ + + +BYTE_UNITS = ('B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB') + + +class Bytes(object): + """ + Supply several methods dealing with bytes. + """ + def __init__(self, bytes): + self.bytes = bytes + + def convert(self, unit=None, multiple=1024): + """ + Convert bytes in give unit. + If `unit` is None, convert bytes in suitable unit. + Convert `multiple` is default to be 1024. + """ + step = 0 + if not unit: + while self.bytes >= multiple and step < len(BYTE_UNITS) - 1: + self.bytes /= multiple + step += 1 + unit = BYTE_UNITS[step] + + else: # convert to specific unit + index_of_unit = BYTE_UNITS.index(unit) + while len(BYTE_UNITS) - 1 > step and index_of_unit != step: + self.bytes /= multiple + step += 1 + return self.bytes, unit diff --git a/tests/test_unit.py b/tests/test_unit.py new file mode 100644 index 0000000..ffc664f --- /dev/null +++ b/tests/test_unit.py @@ -0,0 +1,8 @@ +from pydu.unit import Bytes + + +class TestBytes: + def test_convert(self): + assert Bytes(1024*1024).convert() == (1, 'MB') + assert Bytes(1024*1024).convert(unit='KB') == (1024, 'KB') + assert Bytes(1000).convert(multiple=1000) == (1, 'KB')