subprocess-util/src/p1st/data_units.py

from enum import Enum


def test():
    for num_bytes in [0, 1, 128, 512, 1024, 4096, 75 * 1000, 750 * 1000, 850 * 1000, 1111 * 1000]:
        print(f'{num_bytes} bytes = {DataUnitConverter.to_unit_auto_str(num_bytes)}')


class DataUnitConverter:
    class DataUnit(Enum):
        B = 'B'
        KB = 'KB'
        MB = 'MB'
        GB = 'GB'
        TB = 'TB'

    # Units sorted from small to large.
    _unit_dict: dict[DataUnit, int] = {
        DataUnit.B: 1,
        DataUnit.KB: 1024,
        DataUnit.MB: 1024 * 1024,
        DataUnit.GB: 1024 * 1024 * 1024,
        DataUnit.TB: 1024 * 1024 * 1024 * 1024,
    }

    @classmethod
    def to_bytes(cls, num: int, unit: DataUnit) -> int:
        return num * cls._unit_dict[unit]

    @classmethod
    def to_unit(cls, num_bytes: int, unit: DataUnit = DataUnit.GB) -> float:
        return round(num_bytes / cls._unit_dict[unit], 3)

    @classmethod
    def to_unit_auto(cls, num_bytes: int) -> tuple[float, DataUnit]:
        for unit, factor in cls._unit_dict.items():
            converted = cls.to_unit(num_bytes, unit)
            if converted <= 999:
                return converted, unit

        return num_bytes, cls.DataUnit.B

    @classmethod
    def to_unit_auto_str(cls, num_bytes: int) -> str:
        value, unit = cls.to_unit_auto(num_bytes)
        return f'{value} {unit.value}'


if __name__ == '__main__':
    test()