from collections.abc import Callable, Iterable, Iterator, Mapping from itertools import islice, tee, zip_longest from django.utils.functional import Promise __all__ = [ "BaseChoiceIterator", "BlankChoiceIterator", "CallableChoiceIterator", "flatten_choices", "normalize_choices", ] class BaseChoiceIterator: """Base class for lazy iterators for choices.""" def __eq__(self, other): if isinstance(other, Iterable): return all(a == b for a, b in zip_longest(self, other, fillvalue=object())) return super().__eq__(other) def __getitem__(self, index): if index < 0: # Suboptimally consume whole iterator to handle negative index. return list(self)[index] try: return next(islice(self, index, index + 1)) except StopIteration: raise IndexError("index out of range") from None def __iter__(self): raise NotImplementedError( "BaseChoiceIterator subclasses must implement __iter__()." ) class BlankChoiceIterator(BaseChoiceIterator): """Iterator to lazily inject a blank choice.""" def __init__(self, choices, blank_choice): self.choices = choices self.blank_choice = blank_choice def __iter__(self): choices, other = tee(self.choices) if not any(value in ("", None) for value, _ in flatten_choices(other)): yield from self.blank_choice yield from choices class CallableChoiceIterator(BaseChoiceIterator): """Iterator to lazily normalize choices generated by a callable.""" def __init__(self, func): self.func = func def __iter__(self): yield from normalize_choices(self.func()) def flatten_choices(choices): """Flatten choices by removing nested values.""" for value_or_group, label_or_nested in choices or (): if isinstance(label_or_nested, (list, tuple)): yield from label_or_nested else: yield value_or_group, label_or_nested def normalize_choices(value, *, depth=0): """Normalize choices values consistently for fields and widgets.""" # Avoid circular import when importing django.forms. from django.db.models.enums import ChoicesType match value: case BaseChoiceIterator() | Promise() | bytes() | str(): # Avoid prematurely normalizing iterators that should be lazy. # Because string-like types are iterable, return early to avoid # iterating over them in the guard for the Iterable case below. return value case ChoicesType(): # Choices enumeration helpers already output in canonical form. return value.choices case Mapping() if depth < 2: value = value.items() case Iterator() if depth < 2: # Although Iterator would be handled by the Iterable case below, # the iterator would be consumed prematurely while checking that # its elements are not string-like in the guard, so we handle it # separately. pass case Iterable() if depth < 2 and not any( isinstance(x, (Promise, bytes, str)) for x in value ): # String-like types are iterable, so the guard above ensures that # they're handled by the default case below. pass case Callable() if depth == 0: # If at the top level, wrap callables to be evaluated lazily. return CallableChoiceIterator(value) case Callable() if depth < 2: value = value() case _: return value try: # Recursive call to convert any nested values to a list of 2-tuples. return [(k, normalize_choices(v, depth=depth + 1)) for k, v in value] except (TypeError, ValueError): # Return original value for the system check to raise if it has items # that are not iterable or not 2-tuples: # - TypeError: cannot unpack non-iterable object # - ValueError: values to unpack return value