[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 <me@adamj.eu>
Co-authored-by: Natalia Bidart <124304+nessita@users.noreply.github.com>

Backport of 07fa79ef2bb3e8cace7bd87b292c6c85230eed05 from main

1 files changed, 26 insertions, 0 deletions

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."""