Source code for quantityfield.fields

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.utils import formats
from django.utils.translation import gettext_lazy as _

import datetime
import warnings
from decimal import Decimal
from pint import Quantity
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast

from quantityfield.helper import check_matching_unit_dimension

from .units import ureg
from .widgets import QuantityWidget

DJANGO_JSON_SERIALIZABLE_BASE = Union[
    None, bool, str, int, float, complex, datetime.datetime
]
DJANGO_JSON_SERIALIZABLE = Union[
    Sequence[DJANGO_JSON_SERIALIZABLE_BASE], Dict[str, DJANGO_JSON_SERIALIZABLE_BASE]
]
NUMBER_TYPE = Union[int, float, Decimal]


class QuantityFieldMixin(object):
    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: Optional[List[str]] = 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 = unit_choices

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

        super(QuantityFieldMixin, self).__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,
                )
                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) -> Optional[NUMBER_TYPE]:
        """
        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__(
                "Field '%s' expected a number but got %r." % (self.name, value),
            ) from e

    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) -> Optional[Quantity]:
        if value is None:
            return None
        return self.ureg.Quantity(value * getattr(self.ureg, self.base_units))

    def to_python(self, value) -> Optional[Quantity]:
        if isinstance(value, Quantity):
            return self.fix_unit_registry(value)

        if value is None:
            return None

        to_number = getattr(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

    # TODO: Add tests, understand, add super call if required
    """
    # This code is untested and not documented. It also does not call the super method
    Therefore it is commented out for the moment (even so it is likely required)

    def get_prep_lookup(self, lookup_type, value):

        if lookup_type in ["lt", "gt", "lte", "gte"]:
            if isinstance(value, self.ureg.Quantity):
                v = value.to(self.base_units)
                return v.magnitude
            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(QuantityFieldMixin, self).formfield(**defaults)


class QuantityFormFieldMixin(object):
    """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_fo_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 getattr(widget, "__name__") in classes_to_ignore

        widget = kwargs.get("widget", None)
        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(QuantityFormFieldMixin, self).__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")

        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 DecimalQuantityFormField(QuantityFormFieldMixin, forms.DecimalField):
    to_number_type = Decimal


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

    def __init__(
        self,
        base_units: str,
        *args,
        unit_choices: Optional[List[str]] = 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 get_db_prep_save(self, value, connection) -> Decimal:
        """
        Get Value that shall be saved to database, make sure it is transformed
        """
        converted = self.to_python(value)
        magnitude = self.get_prep_value(converted)
        return connection.ops.adapt_decimalfield_value(
            magnitude, self.max_digits, self.decimal_places
        )