Add some functionalities to `django-model-utils.Choices`.

choices.py

"""Object to store and manipulate lists of choices.

Requires ``django-model-utils``.
"""
from collections.abc import Iterator
from functools import partialmethod

from django import forms
from django.contrib.postgres.fields import ArrayField
from django.db import models
from django.utils.translation import gettext_lazy as _

from model_utils import Choices as ModelUtilsChoices

from .inlistvalidator import InListValidator


class Choices(ModelUtilsChoices):
    """Add some functionalities to ``django-model-utils.Choices``.

    Parameters:
        *choices: The available choices.
        default: Default value.
        groups (dict): Groups of choices (defined as list of "database-value").

    Terminology:
        - key = database value (usually in snake_case)
        - human key = Python identifier (usually in UpperCamelCase)
        - value = human-readable, label

    Initialization: ``Choices(*choices)`` where ``choices`` can be defined as…
        - Single: ``"db-value", …``  # => PythonIdentifier and human-readable = database-value
        - Double: ``("db-value", _("Human-readable")), …``  # => PythonIdentifier = database-value
        - Triple: ``("db-value", "PythonIdentifier", _("Human-readable")), …``
        - Grouped: ``(_("Group name"), Single|Double|Triple), …``

    Protected properties:
        - ``_db_values``: Set of database values.
        - ``_display_map``: Dictionary mapping database value to human-readable.
        - ``_doubles``: List of choices as (database value, human-readable) - can include optgroups.
          Example: ``[("database-value", "Human-readable"), …]``
               or: ``[("Group name", ("database-value", "Human-readable"), …), …]``
        - ``_identifier_map``: Dictionary mapping Python identifier to database value.
          Example: ``{"PythonIdentifier": "database-value", …}``

    Remark: Methods are named as ``get_*()`` so that it has less chance to be the same as an actual
        choice value. So no dictionary-like ``keys()`` and ``values()``.

    Doc: https://django-model-utils.readthedocs.io/en/3.1.2/utilities.html#choices
    Source code: https://github.com/jazzband/django-model-utils/blob/3.1.2/model_utils/choices.py
    """

    def __init__(self, *choices, default=None, groups=None):
        super().__init__(*choices)
        self._default = default
        self._groups = groups
        self._field = None

    def __add__(self, other):
        parent_choices = super().__add__(other)
        return Choices(*parent_choices)

    def bind_field(self, field):
        """Bind this to a model or form field."""
        self._field = field

    def get_by_human_key(self, human_key):
        """Return the key (aka database value) for the given human key (aka Python identifier)."""
        if human_key in self._identifier_map:
            return self._identifier_map[human_key]
        raise KeyError(human_key)

    def get_default(self):
        """Return the default choice (as key/database value)."""
        return self._default

    def get_group(self, key):
        """Return the keys of the given group."""
        if self._groups:
            return self._groups.get(key)
        return None

    def get_groups(self):
        """Return the groups."""
        return self._groups.keys()

    def get_human_key(self, key=None):
        """Return the human key (aka Python identifier) for the given key (aka database value)."""
        if self._field and key is None:
            key = self.get_selected_key()
        for human_key, k in self._identifier_map.items():
            if k == key:
                return human_key
        raise KeyError(key)

    def get_human_keys(self):
        """Return the list of human keys, in original order."""
        return tuple(self._identifier_map.keys())

    def get_keys(self, group=None):
        """Return the list of keys, in original order, optionally for the given group."""
        if group:
            return self.get_group(group)
        return tuple(self._display_map.keys())

    def get_selected_key(self):
        """Return the selected key from the bound field."""
        if self._field:
            return self._field.value_from_object(self._field.model)
        return None

    def get_validator(self):
        """Return the default validator."""
        return InListValidator(self.get_values())

    def get_value(self, key=None):
        """Return the value (aka human-readable) for the given key (aka database value)."""
        if self._field and key is None:
            key = self.get_selected_key()
        return self._display_map[key]

    def get_values(self):
        """Return the list of values (aka human-readable), in original order."""
        return tuple(self._display_map.values())


# MODEL FIELDS =====================================================================================

class ChoiceModelFieldMixin:
    """Model field for ``Choices``."""

    description = _("Choice")

    def __init__(self, choices, **kwargs):
        kwargs["choices"] = choices
        self._set_choices(choices)
        if self.choices_obj:
            kwargs.setdefault("default", self.choices_obj.get_default())
        kwargs.setdefault("db_index", True)
        if "blank" in kwargs:
            kwargs.setdefault("null", kwargs["blank"])
        super().__init__(**kwargs)

    def contribute_to_class(self, cls, name, *args, **kwargs):
        """Add to the model an helper method to get the human key of this field.

        Remark: Only useful if choices are defined as Triple.
        """
        super().contribute_to_class(cls, name, *args, **kwargs)
        setattr(cls, "{}_choices".format(name), self.choices_obj)

    def _set_choices(self, choices):
        if isinstance(choices, Choices):
            self.choices_obj = choices
            default = choices.get_default()
            if default is not None:
                self.default = default
        else:
            self.choices_obj = None

        if isinstance(choices, Iterator):
            choices = list(choices)
        self.choices = choices or []


class ChoiceCharModelField(ChoiceModelFieldMixin, models.CharField):
    """Model field for ``Choices`` stored as string."""

    def __init__(self, choices, **kwargs):
        kwargs.setdefault("max_length", 30)
        kwargs["null"] = False
        super().__init__(choices, **kwargs)
        if self.blank and self.default is None:
            self.default = ""


class ChoiceIntModelField(ChoiceModelFieldMixin, models.PositiveIntegerField):
    """Model field for ``Choices`` stored as positive integer."""

    pass


class ChoiceSmallIntModelField(ChoiceModelFieldMixin, models.PositiveSmallIntegerField):
    """Model field for ``Choices`` stored as positive small integer."""

    pass


# CHOICES ARRAY ====================================================================================

class ArraySelectMultiple(forms.SelectMultiple):
    """Widget for ``ChoiceArrayModelField``.

    Otherwise an empty selection won't be written back to the model.
    Source: https://gist.github.com/danni/f55c4ce19598b2b345ef#gistcomment-2041847
    """

    def value_omitted_from_data(self, data, files, name):
        """Return whether there's data or files for the widget."""
        return False


class ChoiceArrayModelField(ArrayField):
    """Model field for array of ``Choices``.

    Inspired by: https://gist.github.com/danni/f55c4ce19598b2b345ef
    """

    def __init__(self, *args, **kwargs):
        kwargs.setdefault("default", list)
        super().__init__(*args, **kwargs)

    def contribute_to_class(self, cls, name, *args, **kwargs):
        """Add to the model an helper method to get list of display values of this field."""
        super().contribute_to_class(cls, name, *args, **kwargs)
        setattr(cls, "{}_choices".format(name), self.base_field.choices_obj)
        if hasattr(cls, "_get_ARRAYFIELD_display"):  # Not available during migrations.
            setattr(cls, "get_{}_display".format(self.name), partialmethod(cls._get_ARRAYFIELD_display, field=self))

    def formfield(self, **kwargs):
        """Pass the choices from the base field."""
        defaults = {
            "choices": self.base_field.choices,
            "form_class": forms.MultipleChoiceField,
            "widget": ArraySelectMultiple,
        }
        defaults.update(kwargs)
        # Skip our parent's formfield implementation completely as we don't care for it
        # Remark: Cf. ``super(ArrayField, self)``, with ``ArrayField`` and not ``ChoiceArrayModelField``.
        return super(ArrayField, self).formfield(**defaults)

    def to_python(self, value):
        """Convert the value into the correct Python object."""
        value = super().to_python(value)
        if isinstance(value, list):
            value = [self.base_field.to_python(v) for v in value]
        return value


class ChoiceArrayModelMixin:
    """Model mixin to add capabilities related to array of ``Choices`` fields."""

    def _get_ARRAYFIELD_display(self, field):
        """Return the list of display values."""
        choices = dict(field.base_field.flatchoices)
        values = getattr(self, field.attname)
        return [choices.get(value, value) for value in values]

inlistvalidator.py

from django.core.exceptions import ValidationError
from django.utils.deconstruct import deconstructible
from django.utils.translation import gettext_lazy as _


@deconstructible
class InListValidator:
    """Validate that a value is in the list of allowed values."""

    message = _("Wrong value '%(value)s'. Allowed values: %(allowed_values)s.")
    code = "invalid"

    def __init__(self, allowed_values, message=None, code=None):
        self.allowed_values = allowed_values
        if message:
            self.message = message
        if code:
            self.code = code

    def __call__(self, value):
        """Validate the given value."""
        if value not in self.allowed_values:
            raise ValidationError(self.message, code=self.code, params={
                "value": value, "allowed_values": self.get_allowed_values_as_string(),
            })

    def __eq__(self, other):
        return (
            isinstance(other, self.__class__) and
            self.allowed_values == other.allowed_values and
            self.message == other.message and
            self.code == other.code
        )

    def get_allowed_values_as_string(self):
        """Return the list of allowed values as a string."""
        if isinstance(self.allowed_values, [list, tuple, set]):
            return ", ".join([str(v) for v in self.allowed_values])
        if isinstance(self.allowed_values, dict):
            return ", ".join([str(k) for k in self.allowed_values.keys()])
        return self.allowed_values

templatetag.py

from django import template


register = template.Library()


@register.filter(name="human_key")
def human_key_filter(obj, fieldname):
    """
    Return the human key of a ``Choices`` field.

    Usage: ``{{ obj|human_key:"fieldname" }}``
    Example:
        With ``Foo.STATUSES = Choices((1, "OPEN", "Open"), (2, "CLOSED", "Closed"))``
        and ``Foo.status = ChoiceField(Foo.STATUSES)``.
        If ``foo.status = 1`` (i.e. ``Foo.STATUSES.OPEN``), returns ``Open``.
    """
    return obj.get_human_key(fieldname)

usage-examples.py

VEHICLES = Choices(
    ("RedCar", "Red car"),
    ("GreenCar", "Green car"),
    ("RedTruck", "Red truck"),
    default="RedCar",
    groups={
        "red": ["RedCar", "RedTruck],
        "cars": ["RedCar", "GreenCar"],
    },
)

class Foo(models.Model):
    vehicle = ChoiceCharModelField(VEHICLES, verbose_name="Vehicle")
    vehicles = ChoiceArrayModelField(ChoiceCharModelField(VEHICLES), verbose_name="Vehicles")

# In template:
# {{ foo|"vehicle" }}