Compare length of iterator with as few `next` calls as possible. How can I deduplicate this code?

Question

I'm writing a few simple functions that compare the length of an iterator to an integer, returning as early as possible. Right now I have:

class ILengthComparator:
    def __init__(self, iterator: Iterator):
        self.iterator = iterator

    def has_more(self) -> bool:
        try:
            next(self.iterator)
        except StopIteration:
            return False

        return True

    def __eq__(self, length: int) -> bool:
        if length < 0:
            return False

        if not self.has_more():
            return length == 0

        return self == length-1

    def __gt__(self, length: int) -> bool:
        if length < 1:
            return False

        if not self.has_more():
            return True

        return self > length-1

    def __lt__(self, length: int) -> bool:
        if length < 0:
            return True

        if not self.has_more():
            return False

        return self < length - 1

Clearly this class isn't finished (__neq__, __gte__, __lte__), and anything after the first comparison isn't guaranteed to work (as the iterator may be consumed), but I'm wondering: is it is possible to logically simplify this into a single function with an operator function as an argument (ie operator.eq, operator.lt), or at least just reduce the semi-duplicate parts (ie, every function has a comparison with length against a number, and a case for checking if the iterator is empty - each with a rather ad-hoc return value).

(any suggestions don't need to be recursive or use the has_more method, those were just the first thing I reached for when trying to solve this problem)

Pastebin with the above some quick and dirty unit tests: https://pastebin.com/N65Q2qcJ

Answer 1

To know the answer to your comparison to \$N\$, you need to attempt at least:

• \$N\$ next calls if you test with \$\lt\$ (succeeding will result to False) or \$\ge\$ (succeeding will result to True);
• \$N+1\$ next calls in all 4 other cases.

Obviously, if the iterator length is fewer than \$N\$ there will be less calls.

So, in order to simplify things, you could extract the first \$N+1\$ elements of the iterator and count them to compare this count to \$N\$:

• extracting can be simplified with itertools.islice

• counting can be done using either

  • enumerate and dumping the content into a 1-length collections.deque

    from itertools import islice
    from collections import deque
    
    
    def compare_length(iterable, n, operation):
        chunck = islice(iterable, max(0, n+1))
        try:
            (length, _) = deque(enumerate(chunk, 1), maxlen=1)
        except ValueError:
            length = 0
        return operation(length, n)
    

  • sum

    from itertools import islice
    
    
    def compare_length(iterable, n, operation):
        chunck = islice(iterable, max(0, n+1))
        length = sum(1 for _ in chunck)
        return operation(length, n)
    

The former should be more efficient, the latter more readable. As suggested in the question, this function takes the operation to perform as its third argument and compares min(<iterable length>, n+1) to n. This is sufficient in most cases.

---

If you truly want to reduce the amount of next calls needed for \$\lt\$ and \$\ge\$, you can keep your class approach and factorize the min(<iterable length>, n) part into an helper method:

from itertools import islice


class ILengthComparator:
    def __init__(self, iterable):
        self._iterable = iterable

    def _length_up_to(self, n):
        return sum(1 for _ in islice(self._iterable, max(0, n)))

    def __eq__(self, n):
        return self._length_up_to(n+1) == n

    def __ne__(self, n):
        return self._length_up_to(n+1) != n

    def __lt__(self, n):
        return self._length_up_to(n) < n

    def __le__(self, n):
        return self._length_up_to(n+1) <= n

    def __gt__(self, n):
        return self._length_up_to(n+1) > n

    def __ge__(self, n):
        return self._length_up_to(n) >= n

---

Lastly, if you want to avoid consuming the iterator while performing those tests, you can check itertools.tee.