Source code for quantityfield.fields

import datetime
import typing
import warnings
from collections.abc import Callable, Sequence
from decimal import Decimal
from typing import Any, Union, cast

from django import forms
from django.contrib.admin.options import FORMFIELD_FOR_DBFIELD_DEFAULTS
from django.core.exceptions import ValidationError
from django.db import models
from django.db.models.expressions import BaseExpression
from django.utils import formats
from django.utils.translation import gettext_lazy as _

from pint import Quantity

from .helper import check_matching_unit_dimension
from .units import ureg
from .widgets import QuantityWidget

DJANGO_JSON_SERIALIZABLE_BASE = Union[  # noqa: UP007
    None, bool, str, int, float, complex, datetime.datetime
]
DJANGO_JSON_SERIALIZABLE = Union[  # noqa: UP007
    Sequence[DJANGO_JSON_SERIALIZABLE_BASE], dict[str, DJANGO_JSON_SERIALIZABLE_BASE]
]
NUMBER_TYPE = Union[int, float, Decimal]  # noqa: UP007




class QuantityFieldMixin:
    to_number_type: Callable[[Any], NUMBER_TYPE]

    # TODO: Move these stuff into an Protocol or anything
    #       better defining a Mixin
    value_from_object: Callable[[Any], Any]
    name: str
    validate: Callable
    run_validators: Callable

    """A Django Model Field that resolves to a pint Quantity object"""

    def __init__(
        self,
        base_units: str,
        *args,
        unit_choices: typing.Iterable[str] | None = None,
        **kwargs,
    ):
        """
        Create a Quantity field
        :param base_units: Unit description of base unit
        :param unit_choices: If given the possible unit choices with the same
                             dimension like the base_unit
        """
        if not isinstance(base_units, str):
            raise ValueError(
                'QuantityField must be defined with base units, eg: "gram"'
            )

        self.ureg = ureg

        # we do this as a way of raising an exception if some crazy unit was supplied.
        unit = getattr(self.ureg, base_units)  # noqa: F841

        # if we've not hit an exception here, we should be all good
        self.base_units = base_units

        if unit_choices is None:
            self.unit_choices: list[str] = [self.base_units]
        else:
            self.unit_choices = list(unit_choices)
            # The multi widget expects that the base unit is always present as unit
            # choice.
            # Otherwise we would need to handle special cases for no good reason.
            if self.base_units in self.unit_choices:
                self.unit_choices.remove(self.base_units)
            # Base unit has to be the first choice, always as all values are saved as
            # base unit within the database and this would be the first unit shown
            # in the widget
            self.unit_choices = [self.base_units, *self.unit_choices]

        # Check if all unit_choices are valid
        check_matching_unit_dimension(self.ureg, self.base_units, self.unit_choices)

        super().__init__(*args, **kwargs)

    @property
    def units(self) -> str:
        return self.base_units



    def deconstruct(
        self,
    ) -> tuple[
        str,
        str,
        Sequence[DJANGO_JSON_SERIALIZABLE],
        dict[str, DJANGO_JSON_SERIALIZABLE],
    ]:
        """
        Return enough information to recreate the field as a 4-tuple:

         * The name of the field on the model, if contribute_to_class() has
           been run.
         * The import path of the field, including the class:e.g.
           django.db.models.IntegerField This should be the most portable
           version, so less specific may be better.
         * A list of positional arguments.
         * A dict of keyword arguments.

        """
        super_deconstruct = getattr(super(), "deconstruct", None)
        if not callable(super_deconstruct):
            raise NotImplementedError(
                "Tried to use Mixin on a class that has no deconstruct function. "
            )
        name, path, args, kwargs = super_deconstruct()
        kwargs["base_units"] = self.base_units
        kwargs["unit_choices"] = self.unit_choices
        return name, path, args, kwargs




    def fix_unit_registry(self, value: Quantity) -> Quantity:
        """
        Check if the UnitRegistry from settings is used.
        If not try to fix it but give a warning.
        """
        if isinstance(value, Quantity):
            if not isinstance(value, self.ureg.Quantity):
                # Could be fatal if different unit registers are used but we assume
                # the same is used within one project
                # As we warn for this behaviour, we assume that the programmer
                # will fix it and do not include more checks!
                warnings.warn(
                    "Trying to set value from a different unit register for "
                    "quantityfield. "
                    "We assume the naming is equal but best use the same register as"
                    " for creating the quantityfield.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                return value.magnitude * self.ureg(str(value.units))
            else:
                return value
        else:
            raise ValueError(f"Value '{value}' ({type(value)} is not a quantity.")




    def get_prep_value(self, value: Any) -> NUMBER_TYPE | None:
        """
        Perform preliminary non-db specific value checks and conversions.

        Make sure that we compare/use only values without a unit
        """
        # we store the value in the base units defined for this field
        if value is None:
            return None

        if isinstance(value, Quantity):
            quantity = self.fix_unit_registry(value)
            magnitude = quantity.to(self.base_units).magnitude
        else:
            magnitude = value

        try:
            return self.to_number_type(magnitude)
        except (TypeError, ValueError) as e:
            raise e.__class__(
                f"Field '{self.name}' expected a number but got {value!r}.",
            ) from e




    def get_db_prep_value(self, value, connection, prepared=False):
        """
        Convert value to database-compatible format.
        This is called for both save() operations and filter lookups.
        """
        if prepared:
            return value

        # Use get_prep_value to convert Quantity to magnitude
        return self.get_prep_value(value)




    def value_to_string(self, obj) -> str:
        value = self.value_from_object(obj)
        return str(self.get_prep_value(value))




    def from_db_value(self, value: Any, *args, **kwargs) -> Quantity | None:
        if value is None:
            return None
        return self.ureg.Quantity(value, getattr(self.ureg, self.base_units))




    def to_python(self, value) -> Quantity | None:
        if isinstance(value, Quantity):
            return self.fix_unit_registry(value)

        if value is None:
            return None

        to_number = super().to_python
        if not callable(to_number):
            raise NotImplementedError(
                "Mixin not used with a class that has to_python function"
            )

        value = cast(NUMBER_TYPE, to_number(value))

        return self.ureg.Quantity(value, getattr(self.ureg, self.base_units))




    def clean(self, value, model_instance) -> Quantity:
        """
        Convert the value's type and run validation. Validation errors
        from to_python() and validate() are propagated. Return the correct
        value if no error is raised.

        This is a copy from djangos implementation but modified so that validators
        are only checked against the magnitude as otherwise the default database
        validators will not fail because of comparison errors
        """
        value = self.to_python(value)
        check_value = self.get_prep_value(value)
        self.validate(check_value, model_instance)
        self.run_validators(check_value)
        return value




    def formfield(self, **kwargs):
        defaults = {
            "form_class": self.form_field_class,
            "base_units": self.base_units,
            "unit_choices": self.unit_choices,
        }
        defaults.update(kwargs)
        return super().formfield(**defaults)






class QuantityFormFieldMixin:
    """This formfield allows a user to choose which units they
    wish to use to enter a value, but the value is yielded in
    the base_units
    """

    to_number_type: Callable[[Any], NUMBER_TYPE]

    # TODO: Move these stuff into an Protocol or anything
    #       better defining a Mixin
    validate: Callable
    run_validators: Callable
    error_messages: dict[str, str]
    empty_values: Sequence[Any]
    localize: bool

    def __init__(self, *args, **kwargs):
        self.ureg = ureg
        self.base_units = kwargs.pop("base_units", None)
        if self.base_units is None:
            raise ValueError(
                "QuantityFormField requires a base_units kwarg of a "
                "single unit type (eg: grams)"
            )
        self.units = kwargs.pop("unit_choices", [self.base_units])
        if self.base_units not in self.units:
            self.units.append(self.base_units)

        check_matching_unit_dimension(self.ureg, self.base_units, self.units)

        def is_special_admin_widget(widget) -> bool:
            """
            There are some special django admin widgets, defined
            in django/contrib/admin/options.py in the variable
            FORMFIELD_FOR_DBFIELD_DEFAULTS
            The intention for Integer and BigIntegerField is only to
            define the width.

            They are set through a complicated process of the
            modelform_factory setting formfield_callback to
            ModelForm.formfield_to_dbfield

            As they will overwrite our Widget we check for them and
            will ignore them, if they are set as attribute.

            We still will allow subclasses, so the end user has still
            the possibility to use this widget.
            """
            WIDGETS_TO_IGNORE = [
                FORMFIELD_FOR_DBFIELD_DEFAULTS[models.IntegerField],
                FORMFIELD_FOR_DBFIELD_DEFAULTS[models.BigIntegerField],
            ]
            classes_to_ignore = [
                ignored_widget["widget"].__name__
                for ignored_widget in WIDGETS_TO_IGNORE
            ]
            return widget.__name__ in classes_to_ignore

        widget = kwargs.get("widget")
        if widget is None or is_special_admin_widget(widget):
            widget = QuantityWidget(
                base_units=self.base_units, allowed_types=self.units
            )
        kwargs["widget"] = widget
        super().__init__(*args, **kwargs)



    def prepare_value(self, value):
        if isinstance(value, Quantity):
            return value.to(self.base_units)
        else:
            return value




    def clean(self, value):
        """
        General idea, first try to extract the correct number like done in the other
        classes and then follow the same procedure as in the django default field
        """
        if isinstance(value, list) or isinstance(value, tuple):
            val = value[0]
            units = value[1]
        else:
            # If no multi widget is used
            val = value
            units = self.base_units

        if val in self.empty_values:
            # Make sure the correct functions are called also in case of empty values
            self.validate(None)
            self.run_validators(None)
            return None

        if units not in self.units:
            raise ValidationError(_("%(units)s is not a valid choice") % locals())

        if self.localize:
            val = formats.sanitize_separators(value)

        try:
            val = self.to_number_type(val)
        except (ValueError, TypeError):
            raise ValidationError(
                self.error_messages["invalid"], code="invalid"
            ) from None

        val = self.ureg.Quantity(val, getattr(self.ureg, units)).to(self.base_units)
        self.validate(val.magnitude)
        self.run_validators(val.magnitude)
        return val






class QuantityFormField(QuantityFormFieldMixin, forms.FloatField):
    to_number_type = float





class QuantityField(QuantityFieldMixin, models.FloatField):
    form_field_class = QuantityFormField
    to_number_type = float





class IntegerQuantityFormField(QuantityFormFieldMixin, forms.IntegerField):
    to_number_type = int





class IntegerQuantityField(QuantityFieldMixin, models.IntegerField):
    form_field_class = IntegerQuantityFormField
    to_number_type = int





class BigIntegerQuantityField(QuantityFieldMixin, models.BigIntegerField):
    form_field_class = IntegerQuantityFormField
    to_number_type = int





class PositiveIntegerQuantityField(QuantityFieldMixin, models.PositiveIntegerField):
    form_field_class = IntegerQuantityFormField
    to_number_type = int





class DecimalQuantityFormField(QuantityFormFieldMixin, forms.DecimalField):


    def to_number_type(self, x: object) -> Decimal:
        return Decimal(str(x))






class DecimalQuantityField(QuantityFieldMixin, models.DecimalField):
    form_field_class = DecimalQuantityFormField

    def __init__(
        self,
        base_units: str,
        *args,
        unit_choices: list[str] | None = None,
        verbose_name: str = None,
        name: str = None,
        max_digits: int = None,
        decimal_places: int = None,
        **kwargs,
    ):
        # We try to be friendly as default django, if there are missing argument
        # we throw an error early
        if not isinstance(max_digits, int) or not isinstance(decimal_places, int):
            raise ValueError(
                _(
                    "Invalid initialization for DecimalQuantityField! "
                    "We expect max_digits and decimal_places to be set as integers."
                )
            )
        # and we also check the values to be sane
        if decimal_places < 0 or max_digits < 1 or decimal_places > max_digits:
            raise ValueError(
                _(
                    "Invalid initialization for DecimalQuantityField! "
                    "max_digits and decimal_places need to positive and max_digits"
                    "needs to be larger than decimal_places and at least 1. "
                    "So max_digits=%(max_digits)s and "
                    "decimal_plactes=%(decimal_places)s "
                    "are not valid parameters."
                )
                % locals()
            )

        super().__init__(
            base_units,
            *args,
            unit_choices=unit_choices,
            verbose_name=verbose_name,
            name=name,
            max_digits=max_digits,
            decimal_places=decimal_places,
            **kwargs,
        )



    def to_number_type(self, x: object) -> Decimal:
        return Decimal(str(x))




    def get_db_prep_save(self, value, connection) -> Decimal:
        """
        Get Value that shall be saved to database, make sure it is transformed
        """
        if isinstance(value, BaseExpression):
            return value

        converted = self.to_python(value)
        magnitude = self.get_prep_value(converted)
        return connection.ops.adapt_decimalfield_value(
            magnitude, self.max_digits, self.decimal_places
        )