python-compute/compute/utils/units.py
2024-01-13 00:45:30 +03:00

120 lines
3.3 KiB
Python

# This file is part of Compute
#
# Compute is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Compute is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Compute. If not, see <http://www.gnu.org/licenses/>.
"""Tools for data units convertion."""
from collections.abc import Callable
from enum import StrEnum
from compute.exceptions import InvalidDataUnitError
class DataUnit(StrEnum):
"""Data units enumeration."""
BYTES = 'bytes'
KIB = 'KiB'
MIB = 'MiB'
GIB = 'GiB'
TIB = 'TiB'
KB = 'kb'
MB = 'Mb'
GB = 'Gb'
TB = 'Tb'
KBIT = 'kbit'
MBIT = 'Mbit'
GBIT = 'Gbit'
TBIT = 'Tbit'
@classmethod
def _missing_(cls, name: str) -> 'DataUnit':
for member in cls:
if member.name.lower() == name.lower():
return member
return None
def validate_input(*args: str) -> Callable:
"""Validate data units in functions input."""
to_validate = args
def decorator(func: Callable) -> Callable:
def wrapper(*args: float | str, **kwargs: str) -> Callable:
try:
if kwargs:
for arg in to_validate:
unit = kwargs[arg]
DataUnit(unit)
else:
for arg in args[1:]:
unit = arg
DataUnit(unit)
except ValueError as e:
raise InvalidDataUnitError(e, list(DataUnit)) from e
return func(*args, **kwargs)
return wrapper
return decorator
@validate_input('unit')
def to_bytes(value: float, unit: DataUnit = DataUnit.BYTES) -> float:
"""Convert value to bytes."""
unit = DataUnit(unit)
basis = 2 if unit.endswith('iB') else 10
factor = 125 if unit.endswith('bit') else 1
power = {
DataUnit.BYTES: 0,
DataUnit.KIB: 10,
DataUnit.MIB: 20,
DataUnit.GIB: 30,
DataUnit.TIB: 40,
DataUnit.KB: 3,
DataUnit.MB: 6,
DataUnit.GB: 9,
DataUnit.TB: 12,
DataUnit.KBIT: 0,
DataUnit.MBIT: 3,
DataUnit.GBIT: 6,
DataUnit.TBIT: 9,
}
return value * factor * pow(basis, power[unit])
@validate_input('from_unit', 'to_unit')
def convert(value: float, from_unit: DataUnit, to_unit: DataUnit) -> float:
"""Convert units."""
value_in_bits = to_bytes(value, from_unit) * 8
to_unit = DataUnit(to_unit)
basis = 2 if to_unit.endswith('iB') else 10
divisor = 1 if to_unit.endswith('bit') else 8
power = {
DataUnit.BYTES: 0,
DataUnit.KIB: 10,
DataUnit.MIB: 20,
DataUnit.GIB: 30,
DataUnit.TIB: 40,
DataUnit.KB: 3,
DataUnit.MB: 6,
DataUnit.GB: 9,
DataUnit.TB: 12,
DataUnit.KBIT: 3,
DataUnit.MBIT: 6,
DataUnit.GBIT: 9,
DataUnit.TBIT: 12,
}
return value_in_bits / divisor / pow(basis, power[to_unit])