subprocess-util/src/subprocess_util/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:
    """
    A class to convert between data units.

    Example: 1.000.000 B == 1 MB

    The main purpose is to convert a number with many digits
    to a "larger" unit so that the number is shorter
    and more readable.
    """

    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_str(cls, num_bytes: int, unit: DataUnit = DataUnit.GB) -> str:
        value = round(num_bytes / cls._unit_dict[unit], 3)
        return cls.to_str(value, unit)

    @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 cls.to_str(value, unit)

    @classmethod
    def to_str(cls, value: float, unit: DataUnit):
        return f'{value} {unit.value}'


if __name__ == '__main__':
    test()