django-model-utils/model_utils/choices.py

from __future__ import annotations

import copy
from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, overload

T = TypeVar("T")

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator, Sequence

    # The type aliases defined here are evaluated when the django-stubs mypy plugin
    # loads this module, so they must be able to execute under the lowest supported
    # Python VM:
    # - typing.List, typing.Tuple become obsolete in Pyton 3.9
    # - typing.Union becomes obsolete in Pyton 3.10
    from typing import List, Tuple, Union

    from django_stubs_ext import StrOrPromise

    # The type argument 'T' to 'Choices' is the database representation type.
    _Double = Tuple[T, StrOrPromise]
    _Triple = Tuple[T, str, StrOrPromise]
    _Group = Tuple[StrOrPromise, Sequence["_Choice[T]"]]
    _Choice = Union[_Double[T], _Triple[T], _Group[T]]
    # Choices can only be given as a single string if 'T' is 'str'.
    _GroupStr = Tuple[StrOrPromise, Sequence["_ChoiceStr"]]
    _ChoiceStr = Union[str, _Double[str], _Triple[str], _GroupStr]
    # Note that we only accept lists and tuples in groups, not arbitrary sequences.
    # However, annotating it as such causes many problems.

    _DoubleRead = Union[_Double[T], Tuple[StrOrPromise, Iterable["_DoubleRead[T]"]]]
    _DoubleCollector = List[Union[_Double[T], Tuple[StrOrPromise, "_DoubleCollector[T]"]]]
    _TripleCollector = List[Union[_Triple[T], Tuple[StrOrPromise, "_TripleCollector[T]"]]]


class Choices(Generic[T]):
    """
    A class to encapsulate handy functionality for lists of choices
    for a Django model field.

    Each argument to ``Choices`` is a choice, represented as either a
    string, a two-tuple, or a three-tuple.

    If a single string is provided, that string is used as the
    database representation of the choice as well as the
    human-readable presentation.

    If a two-tuple is provided, the first item is used as the database
    representation and the second the human-readable presentation.

    If a triple is provided, the first item is the database
    representation, the second a valid Python identifier that can be
    used as a readable label in code, and the third the human-readable
    presentation. This is most useful when the database representation
    must sacrifice readability for some reason: to achieve a specific
    ordering, to use an integer rather than a character field, etc.

    Regardless of what representation of each choice is originally
    given, when iterated over or indexed into, a ``Choices`` object
    behaves as the standard Django choices list of two-tuples.

    If the triple form is used, the Python identifier names can be
    accessed as attributes on the ``Choices`` object, returning the
    database representation. (If the single or two-tuple forms are
    used and the database representation happens to be a valid Python
    identifier, the database representation itself is available as an
    attribute on the ``Choices`` object, returning itself.)

    Option groups can also be used with ``Choices``; in that case each
    argument is a tuple consisting of the option group name and a list
    of options, where each option in the list is either a string, a
    two-tuple, or a triple as outlined above.

    """

    @overload
    def __init__(self: Choices[str], *choices: _ChoiceStr):
        ...

    @overload
    def __init__(self, *choices: _Choice[T]):
        ...

    def __init__(self, *choices: _ChoiceStr | _Choice[T]):
        # list of choices expanded to triples - can include optgroups
        self._triples: _TripleCollector[T] = []
        # list of choices as (db, human-readable) - can include optgroups
        self._doubles: _DoubleCollector[T] = []
        # dictionary mapping db representation to human-readable
        self._display_map: dict[T, StrOrPromise | list[_Triple[T]]] = {}
        # dictionary mapping Python identifier to db representation
        self._identifier_map: dict[str, T] = {}
        # set of db representations
        self._db_values: set[T] = set()

        self._process(choices)

    def _store(
        self,
        triple: tuple[T, str, StrOrPromise],
        triple_collector: _TripleCollector[T],
        double_collector: _DoubleCollector[T]
    ) -> None:
        self._identifier_map[triple[1]] = triple[0]
        self._display_map[triple[0]] = triple[2]
        self._db_values.add(triple[0])
        triple_collector.append(triple)
        double_collector.append((triple[0], triple[2]))

    def _process(
        self,
        choices: Iterable[_ChoiceStr | _Choice[T]],
        triple_collector: _TripleCollector[T] | None = None,
        double_collector: _DoubleCollector[T] | None = None
    ) -> None:
        if triple_collector is None:
            triple_collector = self._triples
        if double_collector is None:
            double_collector = self._doubles

        def store(c: tuple[Any, str, StrOrPromise]) -> None:
            self._store(c, triple_collector, double_collector)

        for choice in choices:
            # The type inference is not very accurate here:
            # - we lied in the type aliases, stating groups contain an arbitrary Sequence
            #   rather than only list or tuple
            # - there is no way to express that _ChoiceStr is only used when T=str
            # - mypy 1.9.0 doesn't narrow types based on the value of len()
            if isinstance(choice, (list, tuple)):
                if len(choice) == 3:
                    store(choice)
                elif len(choice) == 2:
                    if isinstance(choice[1], (list, tuple)):
                        # option group
                        group_name = choice[0]
                        subchoices = choice[1]
                        tc: _TripleCollector[T] = []
                        triple_collector.append((group_name, tc))
                        dc: _DoubleCollector[T] = []
                        double_collector.append((group_name, dc))
                        self._process(subchoices, tc, dc)
                    else:
                        store((choice[0], cast(str, choice[0]), cast('StrOrPromise', choice[1])))
                else:
                    raise ValueError(
                        "Choices can't take a list of length %s, only 2 or 3"
                        % len(choice)
                    )
            else:
                store((choice, choice, choice))

    def __len__(self) -> int:
        return len(self._doubles)

    def __iter__(self) -> Iterator[_DoubleRead[T]]:
        return iter(self._doubles)

    def __reversed__(self) -> Iterator[_DoubleRead[T]]:
        return reversed(self._doubles)

    def __getattr__(self, attname: str) -> T:
        try:
            return self._identifier_map[attname]
        except KeyError:
            raise AttributeError(attname)

    def __getitem__(self, key: T) -> StrOrPromise | Sequence[_Triple[T]]:
        return self._display_map[key]

    @overload
    def __add__(self: Choices[str], other: Choices[str] | Iterable[_ChoiceStr]) -> Choices[str]:
        ...

    @overload
    def __add__(self, other: Choices[T] | Iterable[_Choice[T]]) -> Choices[T]:
        ...

    def __add__(self, other: Choices[Any] | Iterable[_ChoiceStr | _Choice[Any]]) -> Choices[Any]:
        other_args: list[Any]
        if isinstance(other, self.__class__):
            other_args = other._triples
        else:
            other_args = list(other)
        return Choices(*(self._triples + other_args))

    @overload
    def __radd__(self: Choices[str], other: Iterable[_ChoiceStr]) -> Choices[str]:
        ...

    @overload
    def __radd__(self, other: Iterable[_Choice[T]]) -> Choices[T]:
        ...

    def __radd__(self, other: Iterable[_ChoiceStr] | Iterable[_Choice[T]]) -> Choices[Any]:
        # radd is never called for matching types, so we don't check here
        other_args = list(other)
        # The exact type of 'other' depends on our type argument 'T', which
        # is expressed in the overloading, but lost within this method body.
        return Choices(*(other_args + self._triples))  # type: ignore[arg-type]

    def __eq__(self, other: object) -> bool:
        if isinstance(other, self.__class__):
            return self._triples == other._triples
        return False

    def __repr__(self) -> str:
        return '{}({})'.format(
            self.__class__.__name__,
            ', '.join("%s" % repr(i) for i in self._triples)
        )

    def __contains__(self, item: T) -> bool:
        return item in self._db_values

    def __deepcopy__(self, memo: dict[int, Any] | None) -> Choices[T]:
        args: list[Any] = copy.deepcopy(self._triples, memo)
        return self.__class__(*args)

    def subset(self, *new_identifiers: str) -> Choices[T]:
        identifiers = set(self._identifier_map.keys())

        if not identifiers.issuperset(new_identifiers):
            raise ValueError(
                'The following identifiers are not present: %s' %
                identifiers.symmetric_difference(new_identifiers),
            )

        args: list[Any] = [
            choice for choice in self._triples
            if choice[1] in new_identifiers
        ]
        return self.__class__(*args)