From 711c0547224de00aee39b8720c706ac4977e89fd Mon Sep 17 00:00:00 2001 From: Nick Pope Date: Mon, 16 Oct 2023 19:25:17 +0100 Subject: [5.0.x] Refs #31262 -- Added __eq__() and __getitem__() to BaseChoiceIterator. This makes it easier to work with lazy iterators used for callables, etc. when extracting items or comparing to lists, e.g. during testing. Also added `BaseChoiceIterator.__iter__()` to make it clear that subclasses must implement this and added `__all__` to the module. Co-authored-by: Adam Johnson Co-authored-by: Natalia Bidart <124304+nessita@users.noreply.github.com> Backport of 07fa79ef2bb3e8cace7bd87b292c6c85230eed05 from main --- django/utils/choices.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) (limited to 'django/utils') diff --git a/django/utils/choices.py b/django/utils/choices.py index a0611d96f1..734b9331a1 100644 --- a/django/utils/choices.py +++ b/django/utils/choices.py @@ -1,11 +1,37 @@ from collections.abc import Callable, Iterable, Iterator, Mapping +from itertools import islice, zip_longest from django.utils.functional import Promise +__all__ = [ + "BaseChoiceIterator", + "CallableChoiceIterator", + "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 CallableChoiceIterator(BaseChoiceIterator): """Iterator to lazily normalize choices generated by a callable.""" -- cgit v1.3