Using a condition with `next` on a Python class

Question

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?

Yes, you just need to make `__next__` only return the values you are interested in and not let it return anything else. (If neither of the `if` conditions is true, the function will return `None`, in case you're wondering.) — mkrieger1, Jun 26 '20 at 21:35
@mkrieger1 the `if` condition is already there but still returns `None`... — PedroA, Jun 26 '20 at 21:36
Yes, because if the control flow reaches the end of any function without hitting a `return` statement, the function will return `None`. You need to make sure that there is a `return` statement in all possible cases. — mkrieger1, Jun 26 '20 at 21:37
@mkrieger1 I know, but AFAIU will always return `None` as the default and was wondering if there was some way to avoid that. — PedroA, Jun 26 '20 at 21:41
Yes, by making the control flow reach a `return` statement where you return something different. — mkrieger1, Jun 26 '20 at 21:42
OK, we're going in circles here. The idea is to omit None (or whatever other value I could choose from) from the `__next__` method. I can test for None, or whatever value, in the `for` loop, but "would be nicer to have that class only emit" the values that pass the control flow. That was more my question. — PedroA, Jun 26 '20 at 21:53
Maybe this helps: https://stackoverflow.com/questions/51498464/how-to-prevent-python-function-from-returning-none — mkrieger1, Jun 26 '20 at 21:55
The question is: What should `__next__` return when `self.it` is less than `self._min`. You need to decide that and then write a `return` statement which returns the value you think is appropriate in this case. — mkrieger1, Jun 26 '20 at 21:57
@mkrieger1 the "going in circles" bit actually gave me the idea to test a recursion and it seems to be working. — PedroA, Jun 26 '20 at 22:37

PedroA · Accepted Answer · 2020-06-26T22:36:24.240

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

Using a condition with `__next__` on a Python class

1 Answers1

Using a condition with `next` on a Python class