Implement length-checking iterator

src/sage/misc/misc_c.pyx

@@ -610,3 +610,118 @@ cpdef list normalize_index(object key, int size):
         else:
             raise TypeError("index must be an integer or slice")
     return return_list
+
+
+cdef class sized_iter:
+    """
+    Wrapper for an iterator to verify that it has a specified length.
+
+    INPUT:
+
+    - ``iterable`` -- object to be iterated over
+
+    - ``length`` -- (optional) the required length. If this is not
+      given, then ``len(iterable)`` will be used.
+
+    If the iterable does not have the given length, a ``ValueError`` is
+    raised during iteration.
+
+    EXAMPLES::
+
+        sage: from sage.misc.misc_c import sized_iter
+        sage: list(sized_iter(range(4)))
+        [0, 1, 2, 3]
+        sage: list(sized_iter(range(4), 4))
+        [0, 1, 2, 3]
+        sage: list(sized_iter(range(5), 4))
+        Traceback (most recent call last):
+        ...
+        ValueError: sequence too long (expected length 4, got more)
+        sage: list(sized_iter(range(3), 4))
+        Traceback (most recent call last):
+        ...
+        ValueError: sequence too short (expected length 4, got 3)
+
+    If the iterable is too long, we get the error on the last entry::
+
+        sage: it = sized_iter(range(5), 2)
+        sage: next(it)
+        0
+        sage: next(it)
+        Traceback (most recent call last):
+        ...
+        ValueError: sequence too long (expected length 2, got more)
+
+    When the expected length is zero, the iterator is checked on
+    construction::
+
+        sage: list(sized_iter([], 0))
+        []
+        sage: sized_iter([1], 0)
+        Traceback (most recent call last):
+        ...
+        ValueError: sequence too long (expected length 0, got more)
+
+    If no ``length`` is given, the iterable must implement ``__len__``::
+
+        sage: sized_iter(x for x in range(4))
+        Traceback (most recent call last):
+        ...
+        TypeError: object of type 'generator' has no len()
+    """
+    cdef iterator
+    cdef Py_ssize_t index, size
+
+    def __init__(self, iterable, length=None):
+        self.iterator = iter(iterable)
+        self.index = 0
+        if length is None:
+            self.size = len(iterable)
+        else:
+            self.size = length
+        self.check()
+
+    def __iter__(self):
+        return self
+
+    def __len__(self):
+        """
+        Number of entries remaining, assuming that the expected length
+        is the actual length.
+
+        EXAMPLES::
+
+            sage: from sage.misc.misc_c import sized_iter
+            sage: it = sized_iter(range(4), 4)
+            sage: len(it)
+            4
+            sage: next(it)
+            0
+            sage: len(it)
+            3
+        """
+        return self.size - self.index
+
+    cdef inline int check(self) except -1:
+        """
+        If the iterator is supposed to be exhausted, check that it is.
+        """
+        if self.index < self.size:
+            return 0
+        try:
+            next(self.iterator)
+        except StopIteration:
+            pass
+        else:
+            raise ValueError(f"sequence too long (expected length {self.size}, got more)")
+
+    def __next__(self):
+        if self.index >= self.size:
+            raise StopIteration
+        try:
+            x = next(self.iterator)
+        except StopIteration:
+            raise ValueError(f"sequence too short (expected length {self.size}, got {self.index})")
+        self.index += 1
+        self.check()
+        return x