Search code examples
python-3.xclassiterator

Using a condition with `__next__` on a Python class


Is it possible to implement a condition in a class iterator with the __next__ method? In a case like the one below, if the condition is not met the iterator returns None, but I'd like to omit these and only receive the actual values. For example:

class Foo2:
    def __init__(self, _min=1, _max=5):
        self._max = _max
        self._min = _min
        self.it = -1

    def __iter__(self):
        return self

    def __next__(self):
        self.it += 1
        if self.it == self._max:
            raise StopIteration
        if self.it > self._min:
            return self.it

fu = Foo2(_min=2, _max=5)
for x in fu:
    # if x:  -> can this be controlled by the __next__ method in the class?
    print(x)

# prints
None
None
None
3
4

I'd like to only print the actual values 3 and 4 but instead of testing for None in the loop, would be nicer to have the class only emit those. Is it possible?


Solution

  • I don't know if this is the most correct approach, as there may be problems/drawbacks that I may be overlooking (would actually appreciate if someone could point these out to me), but for this particular case, if the interest is to only ignore values based on some condition, calling next recursively seems to do the trick:

    def __next__(self):
        self.it += 1
        if self.it == self._max:
            raise StopIteration
        elif self.it > self._min:
            return self.it
        else:
            return self.__next__()
    

    Works in a for loop and by calling next directly:

    fu = Foo2(_min=2, _max=5)
    for x in fu:
        print(x)
    
    # prints
    3
    4
    
    fu = Foo2(_min=2, _max=5)
    print(next(fu))  # 3
    print(next(fu))  # 4
    print(next(fu))  # raise StopIteration