import re
from dataclasses import dataclass
from decimal import Decimal
from django.db import models
from ferrosoft.apps.ferrobase.models import MeasurementUnit, UnitRef
from ferrosoft.apps.ferrobase.models.registry import ModelRegistry
from ferrosoft.decimal import zero
[docs]
@dataclass(frozen=True)
class DecimalWithUnit:
"""
DecimalWithUnit wraps a decimal with a unit of measurement.
"""
value: Decimal
unit: MeasurementUnit | None
_GCO2E_PATTERN = re.compile("^(kg|g|[kMG]?t)(CO2e.*)?$")
_MASSES = {
"g": Decimal(1_000),
"kg": Decimal(1_000_000),
"t": Decimal(1_000_000_000),
"kt": Decimal(1_000_000_000_000),
"Mt": Decimal(1_000_000_000_000_000),
"Gt": Decimal(1_000_000_000_000_000_000),
}
_FORMAT_EXP = Decimal("1.00")
[docs]
@classmethod
def create(
cls, value: float | str | Decimal, symbol: UnitRef | None
) -> "DecimalWithUnit":
return cls(
Decimal(value), None if symbol is None else ModelRegistry.Unit.get(symbol)
)
[docs]
@classmethod
def zero(cls, unit: UnitRef | None = None):
return cls.create(zero, unit)
[docs]
def try_unit(self):
"""Return unit or raise exception."""
if self.unit is None:
raise ValueError("Required unit is missing")
return self.unit
[docs]
def with_unit(self, symbol: UnitRef | None) -> "DecimalWithUnit":
return DecimalWithUnit(
self.value, None if symbol is None else ModelRegistry.Unit.get(symbol)
)
[docs]
def is_unit(self, symbol: UnitRef) -> bool:
return self.try_unit().symbol == ModelRegistry.Unit.get(symbol).symbol
[docs]
def round(self, decimal_places: int = 0) -> "DecimalWithUnit":
return DecimalWithUnit(
self.value.quantize(Decimal("1.".ljust(decimal_places + 2, "0"))),
self.unit,
)
[docs]
def symbols_differ(self, other: "DecimalWithUnit") -> bool:
return (
self.unit is not None
and other.unit is not None
and self.unit.symbol != other.unit.symbol
)
def __str__(self):
if self.unit is None:
return str(self.value)
return "%s %s" % (self.value, self.unit.symbol)
def __eq__(self, other):
if isinstance(other, DecimalWithUnit):
value_equals = self.value == other.value
if self.unit is None or other.unit is None:
return value_equals
return value_equals and not self.symbols_differ(other)
if isinstance(other, Decimal):
return self.value == other
raise ValueError("Cannot compare with %s" % type(other))
def __ne__(self, other):
return not self == other
def __lt__(self, other):
if isinstance(other, DecimalWithUnit):
is_lt = self.value < other.value
return is_lt and not self.symbols_differ(other)
elif isinstance(other, Decimal):
return self.value < other
raise ValueError("Cannot < with %s" % type(other))
def __add__(self, other):
if isinstance(other, DecimalWithUnit):
if self.symbols_differ(other):
raise ValueError(
"Cannot add two DecimalWithUnit with different symbols"
)
return DecimalWithUnit(self.value + other.value, self.unit or other.unit)
if isinstance(other, Decimal):
return DecimalWithUnit(self.value + other, self.unit)
raise ValueError("Cannot add with %s" % type(other))
def __sub__(self, other):
if isinstance(other, DecimalWithUnit):
if self.symbols_differ(other):
raise ValueError(
"Cannot subtract two DecimalWithUnit with different symbols"
)
return DecimalWithUnit(self.value - other.value, self.unit or other.unit)
if isinstance(other, Decimal):
return DecimalWithUnit(self.value - other, self.unit)
raise ValueError("Cannot subtract with %s" % type(other))
def __truediv__(self, other):
if isinstance(other, DecimalWithUnit):
if self.symbols_differ(other):
raise ValueError(
"Cannot divide two DecimalWithUnit with different symbols"
)
return DecimalWithUnit(self.value / other.value, self.unit or other.unit)
if isinstance(other, Decimal):
return DecimalWithUnit(self.value / other, self.unit)
raise ValueError("Cannot divide with %s" % type(other))
def __mul__(self, other):
if isinstance(other, DecimalWithUnit):
if self.unit is not None and other.unit is not None:
if self.unit.symbol == other.unit.symbol:
target_unit = self.unit.symbol
else:
target_unit = "%s%s" % (self.unit.symbol, other.unit.symbol)
elif self.unit is not None:
target_unit = self.unit
elif other.unit is not None:
target_unit = other.unit
else:
target_unit = None
return DecimalWithUnit.create(self.value * other.value, target_unit)
if isinstance(other, Decimal):
return DecimalWithUnit(self.value * other, self.unit)
raise ValueError("Cannot multiply with %s" % type(other))
def __pow__(self, other):
if isinstance(other, DecimalWithUnit):
if other.unit is not None:
raise ValueError("Exponent must not have a unit")
return DecimalWithUnit(self.value**other.value, self.unit)
if isinstance(other, Decimal):
return DecimalWithUnit(self.value**other, self.unit)
raise ValueError("Cannot power with %s" % type(other))
def __rmul__(self, other):
return self * other
[docs]
class DecimalWithUnitField(models.DecimalField):
"""A decimal value with a certain presumed unit."""
def __init__(self, *args, unit: MeasurementUnit, **kwargs):
super().__init__(*args, **kwargs)
self.unit = unit
[docs]
def from_db_value(self, value, expression, connection, **kwargs):
if value is None:
return None
return DecimalWithUnit(value, self.unit)
[docs]
def to_python(self, value):
if value is None or isinstance(value, DecimalWithUnit):
return value
return DecimalWithUnit(value, self.unit)