mirror of
https://github.com/Hopiu/django-model-utils.git
synced 2026-03-17 04:10:24 +00:00
233 lines
9.1 KiB
Python
233 lines
9.1 KiB
Python
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)
|